#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <string>

#include "core/JsonHardwareInfo.hpp"
#include "crypto/Crypto.h"
#include "hyper/Hyper.h"
#include "sst/SST.h"

namespace mlir {
#define GEN_PASS_DEF_CONVERTCRYPTOTOARM
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

using namespace mlir;

// SHA-1, MD5, and SM3 output lengths
#define SHA1_DIGEST_LENGTH 20L  // SHA-1 produces a 160-bit (20-byte) hash value
#define MD5_DIGEST_LEN 16L      // MD5 produces a 128-bit (16-byte) hash value
#define SM3_DIGEST_LENGTH 32L   // SM3 produces a 256-bit (32-byte) hash value

enum OP_TYPE {
  INPUT_MULTI_SHA1,
  INPUT_MULTI_MD5,
  INPUT_MULTI_SM3,
  INPUT_MULTI_AES,
  
  INPUT_CRACK_SHA1,
  INPUT_CRACK_MD5,
  INPUT_CRACK_SM3,
  INPUT_CRACK_AES,

  OUTPUT_MULTI_SHA1,
  OUTPUT_MULTI_MD5,
  OUTPUT_MULTI_SM3,
  OUTPUT_MULTI_AES,

  OUTPUT_CRACK_SHA1,
  OUTPUT_CRACK_MD5,
  OUTPUT_CRACK_SM3,
  OUTPUT_CRACK_AES,

  EXTRA_CRACK_SHA1,
  EXTRA_CRACK_MD5,
  EXTRA_CRACK_SM3,
  EXTRA_MULTI_CRACK_AES,
};

void mlirPrintScheduleInfo(PatternRewriter &rewriter, Operation *op,
                           int groupid, Value groupMemIn, Value msglen, int opType) {
  auto loc = op->getLoc();
  wafer::Root &HardwareInfo = wafer::Root::getInstance();
  // create extern function printScheduleInfo if not exist.
  StringRef print_func_name = "printScheduleInfo";
  auto m = op->getParentOfType<ModuleOp>();
  func::FuncOp existingFunc = m.lookupSymbol<func::FuncOp>(print_func_name);
  auto datatype =
      mlir::MemRefType::get(llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                            rewriter.getIntegerType(8));
  auto devicestype =
      mlir::MemRefType::get(llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                            rewriter.getI32Type());
  auto ratiostype =
      mlir::MemRefType::get(llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                            rewriter.getF32Type());
  if (!existingFunc) {
    auto print_schedule_ty = rewriter.getFunctionType(
        mlir::TypeRange{rewriter.getI32Type(),  datatype, devicestype,
                        ratiostype, rewriter.getI32Type(), rewriter.getI32Type()},
        mlir::TypeRange{});
    auto print_schedule =
        func::FuncOp::create(loc, "printScheduleInfo", print_schedule_ty);
    print_schedule.setVisibility(mlir::SymbolTable::Visibility::Private);
    m.push_back(print_schedule);
  }
  // prepare params
  auto groupid_for_print_schedule = rewriter.create<arith::ConstantOp>(
      loc, IntegerAttr::get(rewriter.getI32Type(), groupid));
  Value data_for_print_schedule = groupMemIn;

  float pergroupsum = 0.0;
  auto device_count = 0;
  for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
    if ((c.getType() == "CPU" || c.getType() == "GPU" || c.getType() == "CGRA") &&
        c.getGroup() == groupid) {
      pergroupsum += c.getDutyRatio();
      device_count++;
    }
  }
  auto device_count_value =
      rewriter.create<arith::ConstantIndexOp>(loc, device_count);
  auto devices_for_print_schedule = rewriter.create<memref::AllocaOp>(
      loc, devicestype, ValueRange{device_count_value});
  auto ratios_for_print_schedule = rewriter.create<memref::AllocaOp>(
      loc, ratiostype, ValueRange{device_count_value});
  //  invoke extern function
  int idx = 0;
  for (wafer::Chip c :
       HardwareInfo.getArchitecture().getChips().getChipArray()) {
    if (c.getGroup() != groupid ||
        (c.getType() != "CPU" && c.getType() != "GPU" &&  c.getType() != "CGRA")) {
      continue;
    }
    assert(pergroupsum != 0.0 && "exist lazy group!");
    // create index
    auto idxValue = rewriter.create<arith::ConstantIndexOp>(loc, idx);
    // store ratio
    float tmp_ratio = c.getDutyRatio() / pergroupsum;
    auto ratioValue = rewriter.create<arith::ConstantOp>(
        loc, rewriter.getF32FloatAttr(tmp_ratio));
    rewriter.create<memref::StoreOp>(loc, ratioValue, ratios_for_print_schedule,
                                     ValueRange{idxValue});
    // store device
    Value deviceValue;
    if (c.getType() == "CPU") {
      deviceValue = rewriter.create<arith::ConstantOp>(
          loc, IntegerAttr::get(rewriter.getI32Type(), 0));
    } else if (c.getType() == "GPU") { // GPU
      deviceValue = rewriter.create<arith::ConstantOp>(
          loc, IntegerAttr::get(rewriter.getI32Type(), 1));
    } else if (c.getType() == "CGRA") { // GPU
      deviceValue = rewriter.create<arith::ConstantOp>(
          loc, IntegerAttr::get(rewriter.getI32Type(), 2));
    } else {
      std::cout << "unknown device: " << c.getType() << std::endl;
      assert(false && "unknown device type.");
    }
    rewriter.create<memref::StoreOp>(
        loc, deviceValue, devices_for_print_schedule, ValueRange{idxValue});
    idx += 1;
  }

  auto opTypeI32 = rewriter.create<arith::ConstantOp>(loc, IntegerAttr::get(rewriter.getI32Type(), opType));

  rewriter.create<func::CallOp>(
      loc, print_func_name, mlir::TypeRange{},
      ValueRange{groupid_for_print_schedule,
                 data_for_print_schedule, devices_for_print_schedule,
                 ratios_for_print_schedule, msglen, opTypeI32});
}

extern void convert();

//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//

namespace {
class CryptoSha1MultiToHyperPattern : public OpRewritePattern<crypto::Sha1multi_OP> {
public:
  using OpRewritePattern<crypto::Sha1multi_OP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::Sha1multi_OP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    // get operands of 'crypto.sha1multi'
    Value plaintext = op->getOperand(0);
    Value plainLen = op->getOperand(1);
    Value hash = op->getOperand(2);
    Value num = op->getOperand(3);

    // types
    auto dynI8MemrefType =
        MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    // constants
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, 1);
    Value c20 = rewriter.create<arith::ConstantIndexOp>(loc, 20);
    Value c64 = rewriter.create<arith::ConstantIndexOp>(loc, 64);
    Value c256 = rewriter.create<arith::ConstantIndexOp>(loc, 256);
    // index cast
    Value plainLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/plainLen);
    Value numIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/num);

    Value Nfp =
        rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(), numIdx);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);

    // 'crypto.padmessage.get_padmem'
    auto getPadMemOp = rewriter.create<crypto::PadMessageGetPadmem_OP>(
        loc, /*paddedMem=*/dynI8MemrefType,
        /*paddedLen=*/rewriter.getIndexType(), /*input=*/plaintext,
        /*messageLen=*/plainLen, /*messageNum=*/num);
    Value paddedMem = getPadMemOp.getResult(0);
    Value paddedLen = getPadMemOp.getResult(1);

    // first we need cpuid info
    Value cpu_id = rewriter.create<sst::GetCpuIdOp>(
        loc, /*cpuId=*/rewriter.getIndexType());
    // Value cpu_id = rewriter.create<arith::DivUIOp>(loc,
    // /*lhs=*/cpu_identifier, /*rhs=*/c256);

    wafer::Root &HardwareInfo = wafer::Root::getInstance();

    std::size_t group_num = HardwareInfo.getArchitecture().getChips().CPUNum();

    // if sst use gpu, we use 'RegisterFatbin' and 'RegisterFunc' op.
    bool useGPU = false;
    bool useCGRA = false;
    for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
      if (c.getType() == "GPU" && c.getDutyRatio() > 0) {
        useGPU = true;
      }else if (c.getType() == "CGRA" && c.getDutyRatio() > 0) {
        useCGRA = true;
      }
    }
    if (useGPU) {
      Value handle = rewriter.create<sst::RegisterFatbinOp>(
          loc, /*type=*/rewriter.getIndexType());
      rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle,
                                           /*hostFunc=*/c1);
    }
    if (useCGRA) {
      convert();
    }

    // Task split for multiple CPUs and GPUs
    float sum = 0.0;
    for (size_t i = 0; i < group_num; i++) {
      SmallVector<Attribute> deviceSplit;
      float pergroupsum = 0.0;
      for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if ((c.getType() == "CPU" || c.getType() == "GPU" ||
             c.getType() == "CGRA") &&
            (c.getGroup() == i)) {
          pergroupsum += c.getDutyRatio();
        }
      }

      // when cpuid is right, get attr in hyperfor
      Value sel_cpu = rewriter.create<arith::ConstantIndexOp>(loc, i);
      Value isCpu = rewriter.create<arith::CmpIOp>(
          loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/sel_cpu,
          /*rhs=*/cpu_id);
      rewriter.create<scf::IfOp>(
          loc, /*cond=*/isCpu,
          /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
            Value lowerRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());
            sum += pergroupsum;
            Value upperRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());

            Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
            Value lowerI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), lowerF);
            Value lowerBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), lowerI);

            Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
            Value upperI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), upperF);
            Value upperBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), upperI);
            if (i == group_num - 1) {
              upperBound = numIdx;
            }

            Value groupNumIdx = rewriter.create<arith::SubIOp>(
                loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);
            Value groupSizeIn = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/paddedLen);
            Value groupSizeOut = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/c20);
            Value byte_shift =
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/paddedLen);
            Value groupMemIn = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/paddedMem,
                /*byte_shift=*/byte_shift,
                /*sizes=*/ValueRange{groupSizeIn});
            Value groupMemOut = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/hash,
                /*byte_shift=*/
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/c20),
                /*sizes=*/ValueRange{groupSizeOut});

            auto ENV_printschedule = std::getenv("PRINT_SCHEDULE");
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, plaintext, plainLen, 
                                    static_cast<int>(OP_TYPE::INPUT_MULTI_SHA1));

            // TODO: now we copy all paddedmem, we can only copy necessary
            // paddedmem in the future 'hyper.for' to run multiple no-dependency
            // SHA1 tasks
            auto hyperForOp = rewriter.create<hyper::ForOp>(
                loc, /*lowerBound=*/c0, /*upperBound=*/groupNumIdx,
                /*memIns=*/ValueRange{groupMemIn},
                /*memOuts=*/ValueRange{groupMemOut},
                /*odsArg4=*/
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange miargs, ValueRange moargs) {
                  Value inShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/paddedLen);
                  Value outShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/c20);
                  // extract a single message out
                  Value partMemIn = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/miargs[0], /*byte_shift=*/inShift,
                      /*sizes=*/ValueRange{paddedLen});
                  Value partMemOut = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/moargs[0], /*byte_shift=*/outShift,
                      /*sizes=*/ValueRange{c20});
                  // SHA1 padding
                  rewriter.create<crypto::PadMessagePadone_OP>(
                      loc, /*singleInput=*/partMemIn, /*messageLen=*/plainLen);
                  // SHA1 computation
                  rewriter.create<crypto::Sha1core_OP>(
                      loc, /*singleInput=*/partMemIn,
                      /*singleOutput=*/partMemOut, plainLen);

                  builder.create<hyper::YieldOp>(loc);
                });

            auto ci20 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32IntegerAttr(SHA1_DIGEST_LENGTH));
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemOut, ci20, 
                                    static_cast<int>(OP_TYPE::OUTPUT_MULTI_SHA1));

            for (wafer::Chip c :
                 HardwareInfo.getArchitecture().getChips().getChipArray()) {
              float tmp_ratio;
              if ((c.getType() == "CPU" || c.getType() == "GPU" ||
                   c.getType() == "CGRA") &&
                  c.getGroup() == i) {
                assert(pergroupsum != 0.0 && "exist lazy group!");
                tmp_ratio = c.getDutyRatio() / pergroupsum;
                if (c.getType() == "CPU") {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmp_ratio))}));
                } else if (c.getType() == "CGRA") {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cgra")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmp_ratio))}));
                } else {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("gpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmp_ratio))}));
                }
              }
            }
            // add additional 'hyper.devices' attribute for 'hyper.for' op
            auto deviceSplitAttr = rewriter.getNamedAttr(
                "devices", rewriter.getArrayAttr(deviceSplit));
            SmallVector<NamedAttribute> attrs;
            for (auto attr : hyperForOp->getAttrs()) {
              if (attr.getName() !=
                  "devices") // overwrite default 'hyper.devices' attribute
                attrs.push_back(attr);
            }
            attrs.push_back(deviceSplitAttr);
            hyperForOp->setAttrs(attrs);

            rewriter.create<scf::YieldOp>(loc);
          });
    }

    rewriter.eraseOp(op);

    return success();
  }
};

class Sha1CrackLoweringPattern : public OpRewritePattern<crypto::SHA1CrackOp> {
public:
  using OpRewritePattern<crypto::SHA1CrackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::SHA1CrackOp op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    Value target = op.getTarget();
    Value candidates = op.getCandidates();
    Value numMsg = op.getNumMessages();
    Value msgLen = op.getMessageLen();
    
    // preparations
    // types
    Type dynI8Type =
        MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    Type I8x20Type = MemRefType::get({20}, rewriter.getI8Type());
    // indexes
    Value numMsgIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/numMsg);
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgLen);
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/1);

    Value Nfp = rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(),
                                                    numMsgIdx);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);
    // constants
    Value c0i8 = rewriter.create<arith::ConstantOp>(
        loc, /*value=*/rewriter.getIntegerAttr(rewriter.getI8Type(), 0));
    // memref with initial values to store the crack result
    auto allocOp = rewriter.create<memref::AllocOp>(
        loc, /*memrefType=*/dynI8Type, /*dynamicSizes=*/ValueRange{msgLenIdx});
    allocOp->setAttr("operand_segment_sizes",
                     rewriter.getDenseI32ArrayAttr({1, 0}));
    Value initMemref = allocOp.getResult();
    rewriter.create<scf::ForOp>(
        loc, /*lowerBound=*/c0, /*upperBound=*/msgLenIdx, /*step=*/c1,
        /*iterArgs=*/ValueRange{}, /*bodyBuilder=*/
        [&](OpBuilder &builder, Location loc, Value ivs, ValueRange args) {
          builder.create<memref::StoreOp>(loc, /*valueToStore=*/c0i8,
                                          /*memref=*/initMemref,
                                          /*indices=*/ValueRange{ivs});
          builder.create<scf::YieldOp>(loc);
        });

    // TODO: each cpu in each group just pad theri message(include sha1multi)
    // pad candidate messages outside the loop
    auto getPadMemOp = rewriter.create<crypto::PadMessageGetPadmem_OP>(
        loc, /*paddedMem=*/dynI8Type, /*paddedLen=*/rewriter.getIndexType(),
        /*input=*/candidates, /*messageLen=*/msgLen, /*messageNum=*/numMsg);
    Value paddedCand = getPadMemOp.getResult(0);
    Value paddedLen = getPadMemOp.getResult(1);

    // first we need cpuid info
    Value cpu_id = rewriter.create<sst::GetCpuIdOp>(
        loc, /*cpuId=*/rewriter.getIndexType());

    wafer::Root &HardwareInfo = wafer::Root::getInstance();

    std::size_t group_num = HardwareInfo.getArchitecture().getChips().CPUNum();

    // if sst use gpu, we use 'RegisterFatbin' and 'RegisterFunc' op.
    bool useGPU = false;
    bool useCGRA = false;
    for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
      if (c.getType() == "GPU" && c.getDutyRatio() > 0) {
        useGPU = true;
      }else if (c.getType() == "CGRA" && c.getDutyRatio() > 0) {
        useCGRA = true;
      }
    }
    if (useGPU) {
      Value handle = rewriter.create<sst::RegisterFatbinOp>(
          loc, /*type=*/rewriter.getIndexType());
      rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle,
                                           /*hostFunc=*/c1);
    }
    if (useCGRA) {
      convert();
    }

    // Task split for multiple CPUs and GPUs
    float sum = 0.0;
    for (size_t i = 0; i < group_num; i++) {
      float pergroupsum = 0.0;
      for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if ((c.getType() == "CPU" || c.getType() == "GPU" ||
             c.getType() == "CGRA") &&
            (c.getGroup() == i)) {
          pergroupsum += c.getDutyRatio();
        }
      }

      // when cpuid is right, get attr in hyperfor
      Value sel_cpu = rewriter.create<arith::ConstantIndexOp>(loc, i);
      Value isCpu = rewriter.create<arith::CmpIOp>(
          loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/sel_cpu,
          /*rhs=*/cpu_id);
      rewriter.create<scf::IfOp>(
          loc, /*cond=*/isCpu,
          /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
            Value lowerRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());
            sum += pergroupsum;
            Value upperRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());

            Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
            Value lowerI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), lowerF);
            Value lowerBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), lowerI);

            Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
            Value upperI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), upperF);
            Value upperBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), upperI);
            if (i == group_num - 1) {
              upperBound = numMsgIdx;
            }

            Value groupNumIdx = rewriter.create<arith::SubIOp>(
                loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);
            Value groupSizeIn = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/paddedLen);
            Value byte_shift =
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/paddedLen);
            Value groupMemIn = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8Type, /*source=*/paddedCand,
                /*byte_shift=*/byte_shift,
                /*sizes=*/ValueRange{groupSizeIn});

            auto ENV_printschedule = std::getenv("PRINT_SCHEDULE");
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, candidates, msgLen, 
                                    static_cast<int>(OP_TYPE::INPUT_CRACK_SHA1));

            // use `hyper.for` to calculate SHA1 and compare with target one by
            // one
            auto hyperForOp = rewriter.create<hyper::ForOp>(
                loc, /*lowerBound=*/c0, /*upperBound=*/groupNumIdx,
                /*memIns=*/ValueRange{groupMemIn, target},
                /*memOuts=*/ValueRange{initMemref},
                /*odsArg4=*/
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange miargs, ValueRange moargs) {
                  // data preparation
                  Value idx = builder.create<arith::MulIOp>(loc, /*lhs=*/iv,
                                                            /*rhs=*/paddedLen);
                  // allocate temporary memory for hash result
                  Value alloca = builder.create<memref::AllocaOp>(
                      loc, /*memrefType=*/I8x20Type,
                      /*dynamicSizes=*/ValueRange{});
                  // get original candidate message
                  Value oriMsg = builder.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8Type, /*source=*/miargs[0],
                      /*byte_shift=*/idx, /*sizes=*/ValueRange{msgLenIdx});
                  // get padded memref to be ready for calculation of SHA1
                  Value partMemIn = builder.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8Type, /*source=*/miargs[0],
                      /*byte_shift=*/idx, /*sizes=*/ValueRange{paddedLen});

                  // SHA1 computation
                  builder.create<crypto::PadMessagePadone_OP>(
                      loc, /*singleInput=*/partMemIn, /*messageLen=*/msgLen);
                  builder.create<crypto::Sha1core_OP>(loc, /*input=*/partMemIn,
                                                      /*output=*/alloca, msgLen);

                  // compare with target
                  Value cond = builder.create<crypto::CmpOp>(
                      loc, /*isEqual*/ builder.getI1Type(), /*src=*/alloca,
                      /*dst=*/miargs[1]);
                  builder.create<scf::IfOp>(
                      loc, /*cond=*/cond, /*thenBuilder=*/
                      [&](OpBuilder &builder, Location loc) {
                        // do nothing
                        builder.create<scf::YieldOp>(loc);
                      },
                      /*elseBuilder=*/
                      [&](OpBuilder &builder, Location loc) {
                        // store the result
                        builder.create<scf::ForOp>(
                            loc, /*lowerBound=*/c0, /*upperBound=*/msgLenIdx,
                            /*step=*/c1, /*iterArgs=*/ValueRange{},
                            /*bodyBuilder=*/
                            [&](OpBuilder &builder, Location loc, Value iv,
                                ValueRange args) {
                              // store the result
                              builder.create<memref::StoreOp>(
                                  loc, /*value=*/c0i8, /*memref=*/oriMsg,
                                  /*indices=*/ValueRange{iv});
                              builder.create<scf::YieldOp>(loc);
                            });
                        builder.create<scf::YieldOp>(loc);
                      });

                  // reduce the results
                  builder.create<hyper::ReduceOp>(
                      loc, /*operands=*/ValueRange{oriMsg},
                      /*bodyBuilderFn=*/
                      ArrayRef<mlir::function_ref<void(OpBuilder &, Location,
                                                       Value, Value)>>{
                          [&](OpBuilder &builder, Location loc, Value lhs,
                              Value rhs) {
                            builder.create<scf::ForOp>(
                                loc, /*lowerBound=*/c0,
                                /*upperBound=*/msgLenIdx, /*step=*/c1,
                                /*iterArgs=*/ValueRange{},
                                /*bodyBuilderFns=*/
                                [&](OpBuilder &builder, Location loc, Value iv,
                                    ValueRange args) {
                                  // load one byte
                                  Value val = builder.create<memref::LoadOp>(
                                      loc, /*memref=*/rhs,
                                      /*indices=*/ValueRange{iv});
                                  // atomic "or" operation
                                  builder.create<memref::AtomicRMWOp>(
                                      loc, /*kind=*/
                                      arith::AtomicRMWKindAttr::get(
                                          op.getContext(),
                                          arith::AtomicRMWKind::ori),
                                      /*value=*/val, /*memref=*/lhs,
                                      /*indices=*/ValueRange{iv});
                                  builder.create<scf::YieldOp>(loc);
                                });
                            builder.create<hyper::ReduceReturnOp>(
                                loc, /*result=*/lhs);
                          }});

                  builder.create<hyper::YieldOp>(loc);
                });

            if (ENV_printschedule && std::string(ENV_printschedule) == "ON") {
              mlirPrintScheduleInfo(rewriter, op, i, target, msgLen, 
                                    static_cast<int>(OP_TYPE::EXTRA_CRACK_SHA1));
              mlirPrintScheduleInfo(rewriter, op, i, initMemref, msgLen, 
                                    static_cast<int>(OP_TYPE::OUTPUT_CRACK_SHA1));
            }

            // customized 'devices' attribute for 'hyper.for' op
            SmallVector<Attribute> deviceSplit;
            for (wafer::Chip c :
                 HardwareInfo.getArchitecture().getChips().getChipArray()) {
              float tmp_ratio;
              if ((c.getType() == "CPU" || c.getType() == "GPU" ||
                   c.getType() == "CGRA") &&
                  c.getGroup() == i) {
                assert(pergroupsum != 0.0 && "exist lazy group!");
                tmp_ratio = c.getDutyRatio() / pergroupsum;
                if (c.getType() == "CPU") {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmp_ratio))}));
                } else if (c.getType() == "CGRA") {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cgra")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmp_ratio))}));
                } else {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("gpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmp_ratio))}));
                }
              }
            }
            auto deviceSplitAttr = rewriter.getNamedAttr(
                "devices", rewriter.getArrayAttr(deviceSplit));

            // customized 'isSharedMem' attribute for 'hyper.for' op
            auto isSharedMemAttr = rewriter.getNamedAttr(
                "isSharedMem", rewriter.getBoolArrayAttr({false, true, true}));

            // replace default 'devices' and 'isSharedMem' attributes with
            // customized ones
            SmallVector<NamedAttribute> attrs;
            for (auto attr : hyperForOp->getAttrs()) {
              // overwrite default 'devices' and 'isSharedMem' attribute
              if (attr.getName() != "devices" &&
                  attr.getName() != "isSharedMem")
                attrs.push_back(attr);
            }
            attrs.push_back(deviceSplitAttr);
            attrs.push_back(isSharedMemAttr);
            hyperForOp->setAttrs(attrs);

            rewriter.create<scf::YieldOp>(loc);
          });
    }

    // Determine if a cpu is a master cpu(0)
    Value isMasterCpu = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/cpu_id,
        /*rhs=*/c0);

    // Determine if crack result is null
    Value Result0 = rewriter.create<memref::LoadOp>(loc, /*memref=*/initMemref,
                                                    /*indices=*/ValueRange{c0});
    Value isNullResult = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/Result0,
        /*rhs=*/c0i8);

    // when cpu is a master cpu and crack result is null, cpu get crack result
    Value isMasterAndNull =
        rewriter.create<arith::AndIOp>(loc, isMasterCpu, isNullResult);
    rewriter.create<scf::IfOp>(
        loc, /*cond=*/isMasterAndNull,
        /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
          for (size_t i = 1; i < group_num; i++) {
            Value srcid = rewriter.create<arith::ConstantIndexOp>(loc, i);
            rewriter.create<sst::MsgRecvOp>(loc, /*src=*/srcid, /*dst=*/c0,
                                            /*buffer=*/initMemref,
                                            /*size=*/msgLenIdx);
          }
          rewriter.create<scf::YieldOp>(loc);
        });

    // when cpu is not a master cpu and crack result is not null, cpu send crack
    // result Negate isMasterCpu and isNullResult
    Value notMasterCpu = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::ne, /*lhs=*/cpu_id,
        /*rhs=*/c0);
    Value notNullResult = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::ne, /*lhs=*/Result0,
        /*rhs=*/c0i8);
    Value notMasterAndnotNull =
        rewriter.create<arith::AndIOp>(loc, notMasterCpu, notNullResult);
    rewriter.create<scf::IfOp>(
        loc, /*cond=*/notMasterAndnotNull,
        /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
          rewriter.create<sst::MsgSendOp>(loc, /*src=*/cpu_id, /*dst=*/c0,
                                          /*buffer=*/initMemref,
                                          /*size=*/msgLenIdx);
          rewriter.create<scf::YieldOp>(loc);
        });

    // replace the original operation with the result
    rewriter.replaceOp(op, ValueRange{initMemref});
    // rewriter.eraseOp(op);
    return success();
  }
};

} // namespace

//===----------------------------------------------------------------------===//
// MD5
//===----------------------------------------------------------------------===//
namespace {

class CryptoMd5ToHyperPattern : public OpRewritePattern<crypto::MD5multiOP> {
public:
  using OpRewritePattern<crypto::MD5multiOP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::MD5multiOP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    // get operands of 'crypto.md5multi'
    Value allMsgs = op->getOperand(0);
    Value msgLen = op->getOperand(1);
    Value outputBuffer = op->getOperand(2);
    Value msgCounts = op->getOperand(3);

    // types
    auto dynI8MemrefType =
        MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    // constants
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/1);
    Value c16 = rewriter.create<arith::ConstantIndexOp>(loc, 16);
    // index cast
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgLen);
    Value numIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgCounts);
    Value Nfp =
        rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(), numIdx);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);

    // first we need cpuid info
    Value cpuId = rewriter.create<sst::GetCpuIdOp>(
        loc, /*cpuId=*/rewriter.getIndexType());

    wafer::Root &HardwareInfo = wafer::Root::getInstance();

    std::size_t groupNum = HardwareInfo.getArchitecture().getChips().CPUNum();

    // if sst use gpu, we use 'RegisterFatbin' and 'RegisterFunc' op.
    bool useGPU = false;
    bool useCGRA = false;
    for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
      if (c.getType() == "GPU" && c.getDutyRatio() > 0) {
        useGPU = true;
      }else if (c.getType() == "CGRA" && c.getDutyRatio() > 0) {
        useCGRA = true;
      }
    }
    if (useGPU) {
      Value handle = rewriter.create<sst::RegisterFatbinOp>(
          loc, /*type=*/rewriter.getIndexType());
      rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle,
                                           /*hostFunc=*/c1);
    }
    if (useCGRA) {
      convert();
    }

    float sum = 0.0;
    for (size_t i = 0; i < groupNum; i++) {
      SmallVector<Attribute> deviceSplit;
      float perGroupSum = 0.0;
      for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if ((c.getType() == "CPU" || c.getType() == "GPU" ||
             c.getType() == "CGRA") &&
            c.getGroup() == i) {
          perGroupSum += c.getDutyRatio();
        }
      }

      Value selCPU = rewriter.create<arith::ConstantIndexOp>(loc, i);
      Value isCPU = rewriter.create<arith::CmpIOp>(
          loc, arith::CmpIPredicate::eq, selCPU, cpuId);
      rewriter.create<scf::IfOp>(
          loc, isCPU, [&](OpBuilder &builder, Location loc) {
            Value lowerRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());
            sum += perGroupSum;
            Value upperRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());

            Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
            Value lowerI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), lowerF);
            Value lowerBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), lowerI);

            Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
            Value upperI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), upperF);
            Value upperBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), upperI);
            if (i == groupNum - 1) {
              upperBound = numIdx;
            }

            Value groupNumIdx = rewriter.create<arith::SubIOp>(
                loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);
            Value groupSizeIn = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/msgLenIdx);
            Value groupSizeOut = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/c16);
            Value byte_shift =
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/msgLenIdx);
            Value groupMemIn = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/allMsgs,
                /*byte_shift=*/
                byte_shift,
                /*sizes=*/ValueRange{groupSizeIn});
            Value groupMemOut = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/outputBuffer,
                /*byte_shift=*/
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/c16),
                /*sizes=*/ValueRange{groupSizeOut});

            auto ENV_printschedule = std::getenv("PRINT_SCHEDULE");
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemIn, msgLen, 
                                    static_cast<int>(OP_TYPE::INPUT_MULTI_MD5));

            // 'hyper.for' to run multiple no-dependency MD5 tasks
            auto hyperForOp = rewriter.create<hyper::ForOp>(
                loc, /*lowerBound=*/c0, /*upperBound=*/groupNumIdx,
                /*memIns=*/ValueRange{groupMemIn},
                /*memOuts=*/ValueRange{groupMemOut},
                /*odsArg4=*/
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange miargs, ValueRange moargs) {
                  Value inShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/msgLenIdx);
                  Value outShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/c16);
                  // extract a single message out
                  Value input = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/miargs[0], /*byte_shift=*/inShift,
                      /*sizes=*/ValueRange{msgLenIdx});
                  Value output = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/moargs[0], /*byte_shift=*/outShift,
                      /*sizes=*/ValueRange{c16});

                  // MD5 computation
                  rewriter.create<crypto::MD5OP>(loc, /*input=*/input,
                                                 /*msgLen=*/msgLen,
                                                 /*output=*/output);

                  builder.create<hyper::YieldOp>(loc);
                });

            auto ci16 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32IntegerAttr(MD5_DIGEST_LEN));
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemOut, ci16, 
                                    static_cast<int>(OP_TYPE::OUTPUT_MULTI_MD5));

            for (wafer::Chip c :
                 HardwareInfo.getArchitecture().getChips().getChipArray()) {
              float tmpRatio;
              if ((c.getType() == "CPU" || c.getType() == "GPU" || c.getType() == "CGRA") &&
                  c.getGroup() == i) {
                assert(perGroupSum != 0.0 && "exist lazy group!");
                tmpRatio = c.getDutyRatio() / perGroupSum;
                if (c.getType() == "CPU") {
                  // cpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else if (c.getType() == "CGRA") {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cgra")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else {
                  // gpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("gpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                }
              }
            } // end of chip split for-loop

            // add additional 'hyper.devices' attribute for 'hyper.for' op
            auto deviceSplitAttr = rewriter.getNamedAttr(
                "devices", rewriter.getArrayAttr(deviceSplit));
            SmallVector<NamedAttribute> attrs;
            for (auto attr : hyperForOp->getAttrs()) {
              if (attr.getName() !=
                  "devices") // overwrite default 'hyper.devices' attribute
                attrs.push_back(attr);
            }
            attrs.push_back(deviceSplitAttr);
            hyperForOp->setAttrs(attrs);

            rewriter.create<scf::YieldOp>(loc);
          });
    }

    // erase the original 'crypto.md5multi' op
    rewriter.eraseOp(op);

    return success();
  }
};

class MD5CrackLoweringPattern : public OpRewritePattern<crypto::MD5CrackOp> {
public:
  using OpRewritePattern<crypto::MD5CrackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::MD5CrackOp op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    Value target = op.getTarget();
    Value candidates = op.getCandidates();
    Value numMsg = op.getNumMessages();
    Value msgLen = op.getMessageLen();

    // types
    Type dynI8Type =
        MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    Type I8x16Type = MemRefType::get({16}, rewriter.getI8Type());
    // index cast
    Value numIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/numMsg);
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgLen);
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/1);
    Value Nfp =
        rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(), numIdx);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);

    // constants
    Value c0i8 = rewriter.create<arith::ConstantOp>(
        loc, /*value=*/rewriter.getIntegerAttr(rewriter.getI8Type(), 0));
    // memref with initial values to store the crack result
    auto allocOp = rewriter.create<memref::AllocOp>(
        loc, /*memrefType=*/dynI8Type, /*dynamicSizes=*/ValueRange{msgLenIdx});
    allocOp->setAttr("operand_segment_sizes",
                     rewriter.getDenseI32ArrayAttr({1, 0}));
    Value initMemref = allocOp.getResult();
    rewriter.create<scf::ForOp>(
        loc, /*lowerBound=*/c0, /*upperBound=*/msgLenIdx, /*step=*/c1,
        /*iterArgs=*/ValueRange{}, /*bodyBuilder=*/
        [&](OpBuilder &builder, Location loc, Value ivs, ValueRange args) {
          builder.create<memref::StoreOp>(loc, /*valueToStore=*/c0i8,
                                          /*memref=*/initMemref,
                                          /*indices=*/ValueRange{ivs});
          builder.create<scf::YieldOp>(loc);
        });

    // first we need cpuid info
    Value cpuId = rewriter.create<sst::GetCpuIdOp>(
        loc, /*cpuId=*/rewriter.getIndexType());

    wafer::Root &HardwareInfo = wafer::Root::getInstance();

    std::size_t groupNum = HardwareInfo.getArchitecture().getChips().CPUNum();

    // if sst use gpu, we use 'RegisterFatbin' and 'RegisterFunc' op.
    bool useGPU = false;
    bool useCGRA = false;
    for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
      if (c.getType() == "GPU" && c.getDutyRatio() > 0) {
        useGPU = true;
      }else if (c.getType() == "CGRA" && c.getDutyRatio() > 0) {
        useCGRA = true;
      }
    }
    if (useGPU) {
      Value handle = rewriter.create<sst::RegisterFatbinOp>(
          loc, /*type=*/rewriter.getIndexType());
      rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle,
                                           /*hostFunc=*/c1);
    }
    if (useCGRA) {
      convert();
    }

    float sum = 0.0;
    for (size_t i = 0; i < groupNum; i++) {
      float perGroupSum = 0.0;
      for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if ((c.getType() == "CPU" || c.getType() == "GPU" ||
             c.getType() == "CGRA") &&
            c.getGroup() == i) {
          perGroupSum += c.getDutyRatio();
        }
      }

      Value selCPU = rewriter.create<arith::ConstantIndexOp>(loc, i);
      Value isCPU = rewriter.create<arith::CmpIOp>(
          loc, arith::CmpIPredicate::eq, selCPU, cpuId);
      rewriter.create<scf::IfOp>(
          loc, isCPU, [&](OpBuilder &builder, Location loc) {
            Value lowerRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());
            sum += perGroupSum;
            Value upperRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());

            Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
            Value lowerI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), lowerF);
            Value lowerBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), lowerI);

            Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
            Value upperI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), upperF);
            Value upperBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), upperI);
            if (i == groupNum - 1) {
              upperBound = numIdx;
            }

            Value groupNumIdx = rewriter.create<arith::SubIOp>(
                loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);
            Value groupSizeIn = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/msgLenIdx);
            Value byte_shift =
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/msgLenIdx);
            Value groupMemIn = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8Type, /*source=*/candidates,
                /*byte_shift=*/byte_shift,
                /*sizes=*/ValueRange{groupSizeIn});

            auto ENV_printschedule = std::getenv("PRINT_SCHEDULE");
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemIn, msgLen, 
                                    static_cast<int>(OP_TYPE::INPUT_CRACK_MD5));

            // 'hyper.for' to run multiple no-dependency MD5 tasks
            auto hyperForOp = rewriter.create<hyper::ForOp>(
                loc, /*lowerBound=*/c0, /*upperBound=*/groupNumIdx,
                /*memIns=*/ValueRange{groupMemIn, target},
                /*memOuts=*/ValueRange{initMemref},
                /*odsArg4=*/
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange miargs, ValueRange moargs) {
                  Value idx = builder.create<arith::MulIOp>(loc, /*lhs=*/iv,
                                                            /*rhs=*/msgLenIdx);
                  // allocate temporary memory for hash result
                  Value alloca = builder.create<memref::AllocaOp>(
                      loc, /*memrefType=*/I8x16Type,
                      /*dynamicSizes=*/ValueRange{});
                  // get original candidate message
                  Value oriMsg = builder.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8Type, /*source=*/miargs[0],
                      /*byte_shift=*/idx, /*sizes=*/ValueRange{msgLenIdx});
                  // get memref to be ready for calculation of MD5
                  Value partMemIn = builder.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8Type, /*source=*/miargs[0],
                      /*byte_shift=*/idx, /*sizes=*/ValueRange{msgLenIdx});

                  // MD5 computation
                  rewriter.create<crypto::MD5OP>(loc, /*input=*/partMemIn,
                                                 /*msgLen=*/msgLen,
                                                 /*output=*/alloca);

                  // compare with target
                  Value cond = builder.create<crypto::CmpOp>(
                      loc, /*isEqual*/ builder.getI1Type(), /*src=*/alloca,
                      /*dst=*/miargs[1]);
                  builder.create<scf::IfOp>(
                      loc, /*cond=*/cond, /*thenBuilder=*/
                      [&](OpBuilder &builder, Location loc) {
                        // do nothing
                        builder.create<scf::YieldOp>(loc);
                      },
                      /*elseBuilder=*/
                      [&](OpBuilder &builder, Location loc) {
                        // store the result
                        builder.create<scf::ForOp>(
                            loc, /*lowerBound=*/c0, /*upperBound=*/msgLenIdx,
                            /*step=*/c1, /*iterArgs=*/ValueRange{},
                            /*bodyBuilder=*/
                            [&](OpBuilder &builder, Location loc, Value iv,
                                ValueRange args) {
                              // store the result
                              builder.create<memref::StoreOp>(
                                  loc, /*value=*/c0i8, /*memref=*/oriMsg,
                                  /*indices=*/ValueRange{iv});
                              builder.create<scf::YieldOp>(loc);
                            });
                        builder.create<scf::YieldOp>(loc);
                      });

                  // reduce the results
                  builder.create<hyper::ReduceOp>(
                      loc, /*operands=*/ValueRange{oriMsg},
                      /*bodyBuilderFn=*/
                      ArrayRef<mlir::function_ref<void(OpBuilder &, Location,
                                                       Value, Value)>>{
                          [&](OpBuilder &builder, Location loc, Value lhs,
                              Value rhs) {
                            builder.create<scf::ForOp>(
                                loc, /*lowerBound=*/c0,
                                /*upperBound=*/msgLenIdx, /*step=*/c1,
                                /*iterArgs=*/ValueRange{},
                                /*bodyBuilderFns=*/
                                [&](OpBuilder &builder, Location loc, Value iv,
                                    ValueRange args) {
                                  // load one byte
                                  Value val = builder.create<memref::LoadOp>(
                                      loc, /*memref=*/rhs,
                                      /*indices=*/ValueRange{iv});
                                  // atomic "or" operation
                                  builder.create<memref::AtomicRMWOp>(
                                      loc, /*kind=*/
                                      arith::AtomicRMWKindAttr::get(
                                          op.getContext(),
                                          arith::AtomicRMWKind::ori),
                                      /*value=*/val, /*memref=*/lhs,
                                      /*indices=*/ValueRange{iv});
                                  builder.create<scf::YieldOp>(loc);
                                });
                            builder.create<hyper::ReduceReturnOp>(
                                loc, /*result=*/lhs);
                          }});

                  builder.create<hyper::YieldOp>(loc);
                });

            if (ENV_printschedule && std::string(ENV_printschedule) == "ON") {
              mlirPrintScheduleInfo(rewriter, op, i, target, msgLen, 
                                    static_cast<int>(OP_TYPE::EXTRA_CRACK_MD5));
              mlirPrintScheduleInfo(rewriter, op, i, initMemref, msgLen, 
                                    static_cast<int>(OP_TYPE::OUTPUT_CRACK_MD5));
            }
            SmallVector<Attribute> deviceSplit;
            for (wafer::Chip c :
                 HardwareInfo.getArchitecture().getChips().getChipArray()) {
              float tmpRatio;
              if ((c.getType() == "CPU" || c.getType() == "GPU" || c.getType() == "CGRA") &&
                  c.getGroup() == i) {
                assert(perGroupSum != 0.0 && "exist lazy group!");
                tmpRatio = c.getDutyRatio() / perGroupSum;
                if (c.getType() == "CPU") {
                  // cpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else if (c.getType() == "CGRA") {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cgra")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else {
                  // gpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("gpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                }
              }
            } // end of chip split for-loop

            // add additional 'hyper.devices' attribute for 'hyper.for' op
            auto deviceSplitAttr = rewriter.getNamedAttr(
                "devices", rewriter.getArrayAttr(deviceSplit));
            // customized 'isSharedMem' attribute for 'hyper.for' op
            auto isSharedMemAttr = rewriter.getNamedAttr(
                "isSharedMem", rewriter.getBoolArrayAttr({false, true, true}));

            SmallVector<NamedAttribute> attrs;
            for (auto attr : hyperForOp->getAttrs()) {
              if (attr.getName() != "devices" &&
                  attr.getName() != "isSharedMem")
                attrs.push_back(attr);
            }
            attrs.push_back(deviceSplitAttr);
            attrs.push_back(isSharedMemAttr);
            hyperForOp->setAttrs(attrs);

            rewriter.create<scf::YieldOp>(loc);
          });
    }

    // Determine if a cpu is a master cpu(0)
    Value isMasterCpu = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/cpuId, /*rhs=*/c0);

    // Determine if crack result is null
    Value Result0 = rewriter.create<memref::LoadOp>(loc, /*memref=*/initMemref,
                                                    /*indices=*/ValueRange{c0});
    Value isNullResult = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/Result0,
        /*rhs=*/c0i8);

    // when cpu is a master cpu and crack result is null, cpu get crack result
    Value isMasterAndNull =
        rewriter.create<arith::AndIOp>(loc, isMasterCpu, isNullResult);
    rewriter.create<scf::IfOp>(
        loc, /*cond=*/isMasterAndNull,
        /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
          for (size_t i = 1; i < groupNum; i++) {
            Value srcid = rewriter.create<arith::ConstantIndexOp>(loc, i);
            rewriter.create<sst::MsgRecvOp>(loc, /*src=*/srcid, /*dst=*/c0,
                                            /*buffer=*/initMemref,
                                            /*size=*/msgLenIdx);
          }
          rewriter.create<scf::YieldOp>(loc);
        });

    // when cpu is not a master cpu and crack result is not null, cpu send crack
    // result Negate isMasterCpu and isNullResult
    Value notMasterCpu = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::ne, /*lhs=*/cpuId, /*rhs=*/c0);
    Value notNullResult = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::ne, /*lhs=*/Result0,
        /*rhs=*/c0i8);
    Value notMasterAndnotNull =
        rewriter.create<arith::AndIOp>(loc, notMasterCpu, notNullResult);
    rewriter.create<scf::IfOp>(
        loc, /*cond=*/notMasterAndnotNull,
        /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
          rewriter.create<sst::MsgSendOp>(loc, /*src=*/cpuId, /*dst=*/c0,
                                          /*buffer=*/initMemref,
                                          /*size=*/msgLenIdx);
          rewriter.create<scf::YieldOp>(loc);
        });

    // replace the original operation with the result
    rewriter.replaceOp(op, ValueRange{initMemref});
    // rewriter.eraseOp(op);
    return success();
  }

}; // end of MD5Crack

} // namespace

//===----------------------------------------------------------------------===//
// SM3
//===----------------------------------------------------------------------===//
namespace {

class CryptoSm3ToHyperPattern : public OpRewritePattern<crypto::SM3multiOP> {
public:
  using OpRewritePattern<crypto::SM3multiOP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::SM3multiOP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    // get operands of 'crypto.sm3multi'
    Value allInputs = op->getOperand(0);
    Value msgLen = op->getOperand(1);
    Value outputBuffer = op->getOperand(2);
    Value msgCount = op->getOperand(3);

    // types
    auto dynI8MemrefType =
        MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    // constants
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, 1);
    Value c32 = rewriter.create<arith::ConstantIndexOp>(loc, 32);
    // index cast
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgLen);
    Value numIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgCount);
    Value Nfp =
        rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(), numIdx);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);

    // first we need cpuid info
    Value cpuId = rewriter.create<sst::GetCpuIdOp>(
        loc, /*cpuId=*/rewriter.getIndexType());

    wafer::Root &HardwareInfo = wafer::Root::getInstance();

    std::size_t groupNum = HardwareInfo.getArchitecture().getChips().CPUNum();

    // if sst use gpu, we use 'RegisterFatbin' and 'RegisterFunc' op.
    bool useGPU = false;
    bool useCGRA = false;
    for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
      if (c.getType() == "GPU" && c.getDutyRatio() > 0) {
        useGPU = true;
      }else if (c.getType() == "CGRA" && c.getDutyRatio() > 0) {
        useCGRA = true;
      }
    }
    if (useGPU) {
      Value handle = rewriter.create<sst::RegisterFatbinOp>(
          loc, /*type=*/rewriter.getIndexType());
      rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle,
                                           /*hostFunc=*/c1);
    }
    if (useCGRA) {
      convert();
    }

    float sum = 0.0;
    for (size_t i = 0; i < groupNum; i++) {
      SmallVector<Attribute> deviceSplit;
      float perGroupSum = 0.0;
      for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if ((c.getType() == "CPU" || c.getType() == "GPU" ||
             c.getType() == "CGRA") &&
            c.getGroup() == i) {
          perGroupSum += c.getDutyRatio();
        }
      }

      Value selCPU = rewriter.create<arith::ConstantIndexOp>(loc, i);
      Value isCPU = rewriter.create<arith::CmpIOp>(
          loc, arith::CmpIPredicate::eq, selCPU, cpuId);
      rewriter.create<scf::IfOp>(
          loc, isCPU,
          [&](OpBuilder &builder, Location loc) {
            Value lowerRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());
            sum += perGroupSum;
            Value upperRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());

            Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
            Value lowerI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), lowerF);
            Value lowerBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), lowerI);

            Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
            Value upperI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), upperF);
            Value upperBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), upperI);
            if (i == groupNum - 1) {
              upperBound = numIdx;
            }

            Value groupNumIdx = rewriter.create<arith::SubIOp>(
                loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);
            Value groupSizeIn = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/msgLenIdx);
            Value groupSizeOut = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/c32);
            Value byte_shift =
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/msgLenIdx);
            Value groupMemIn = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/allInputs,
                /*byte_shift=*/byte_shift,
                /*sizes=*/ValueRange{groupSizeIn});
            Value groupMemOut = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/outputBuffer,
                /*byte_shift=*/
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/c32),
                /*sizes=*/ValueRange{groupSizeOut});

            auto ENV_printschedule = std::getenv("PRINT_SCHEDULE");
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemIn, msgLen, 
                                    static_cast<int>(OP_TYPE::INPUT_MULTI_SM3));

            auto hyperForOp = rewriter.create<hyper::ForOp>(
                loc, /*lowerBound=*/c0, /*upperBound=*/groupNumIdx,
                /*memIns=*/ValueRange{groupMemIn},
                /*memOuts=*/ValueRange{groupMemOut},
                /*odsArg4=*/
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange miargs, ValueRange moargs) {
                  Value inShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/msgLenIdx);
                  Value outShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/c32);
                  // extract a single message out
                  Value input = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/miargs[0], /*byte_shift=*/inShift,
                      /*sizes=*/ValueRange{msgLenIdx});
                  Value output = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/moargs[0], /*byte_shift=*/outShift,
                      /*sizes=*/ValueRange{c32});

                  rewriter.create<crypto::SM3OP>(loc, /*input=*/input,
                                                 /*msgLen=*/msgLen,
                                                 /*output=*/output);

                  builder.create<hyper::YieldOp>(loc);
                });

            auto ci32 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI32IntegerAttr(SM3_DIGEST_LENGTH));
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemOut, ci32, 
                                    static_cast<int>(OP_TYPE::OUTPUT_MULTI_SM3));

            for (wafer::Chip c :
                 HardwareInfo.getArchitecture().getChips().getChipArray()) {
              float tmpRatio;
              if ((c.getType() == "CPU" || c.getType() == "GPU" || c.getType() == "CGRA") &&
                  c.getGroup() == i) {
                assert(perGroupSum != 0.0 && "exist lazy group!");
                tmpRatio = c.getDutyRatio() / perGroupSum;
                if (c.getType() == "CPU") {
                  // cpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else if (c.getType() == "CGRA") {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cgra")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else {
                  // gpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("gpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                }
              }
            }

            // add additional 'hyper.devices' attribute for 'hyper.for' op
            auto deviceSplitAttr = rewriter.getNamedAttr(
                "devices", rewriter.getArrayAttr(deviceSplit));
            SmallVector<NamedAttribute> attrs;
            for (auto attr : hyperForOp->getAttrs()) {
              if (attr.getName() !=
                  "devices") // overwrite default 'hyper.devices' attribute
                attrs.push_back(attr);
            }
            attrs.push_back(deviceSplitAttr);
            hyperForOp->setAttrs(attrs);

            rewriter.create<scf::YieldOp>(loc);
          } // end of if body
      ); // end of scf::If
    } // end of group loop

    rewriter.eraseOp(op);

    return success();
  }
};

class SM3CrackLoweringPattern : public OpRewritePattern<crypto::SM3CrackOp> {
public:
  using OpRewritePattern<crypto::SM3CrackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::SM3CrackOp op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    Value target = op.getTarget();
    Value candidates = op.getCandidates();
    Value numMsg = op.getNumMessages();
    Value msgLen = op.getMessageLen();
    // types
    Type dynI8Type =
        MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    Type I8x32Type = MemRefType::get({32}, rewriter.getI8Type());
    // index cast
    Value numIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/numMsg);
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgLen);
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/1);
    Value Nfp =
        rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(), numIdx);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);

    // constants
    Value c0i8 = rewriter.create<arith::ConstantOp>(
        loc, /*value=*/rewriter.getIntegerAttr(rewriter.getI8Type(), 0));
    // memref with initial values to store the crack result
    auto allocOp = rewriter.create<memref::AllocOp>(
        loc, /*memrefType=*/dynI8Type, /*dynamicSizes=*/ValueRange{msgLenIdx});
    allocOp->setAttr("operand_segment_sizes",
                     rewriter.getDenseI32ArrayAttr({1, 0}));
    Value initMemref = allocOp.getResult();
    rewriter.create<scf::ForOp>(
        loc, /*lowerBound=*/c0, /*upperBound=*/msgLenIdx, /*step=*/c1,
        /*iterArgs=*/ValueRange{}, /*bodyBuilder=*/
        [&](OpBuilder &builder, Location loc, Value ivs, ValueRange args) {
          builder.create<memref::StoreOp>(loc, /*valueToStore=*/c0i8,
                                          /*memref=*/initMemref,
                                          /*indices=*/ValueRange{ivs});
          builder.create<scf::YieldOp>(loc);
        });

    // first we need cpuid info
    Value cpuId = rewriter.create<sst::GetCpuIdOp>(
        loc, /*cpuId=*/rewriter.getIndexType());

    wafer::Root &HardwareInfo = wafer::Root::getInstance();

    std::size_t groupNum = HardwareInfo.getArchitecture().getChips().CPUNum();

    // if sst use gpu, we use 'RegisterFatbin' and 'RegisterFunc' op.
    bool useGPU = false;
    bool useCGRA = false;
    for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
      if (c.getType() == "GPU" && c.getDutyRatio() > 0) {
        useGPU = true;
      }else if (c.getType() == "CGRA" && c.getDutyRatio() > 0) {
        useCGRA = true;
      }
    }
    if (useGPU) {
      Value handle = rewriter.create<sst::RegisterFatbinOp>(
          loc, /*type=*/rewriter.getIndexType());
      rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle,
                                           /*hostFunc=*/c1);
    }
    if (useCGRA) {
      convert();
    }

    float sum = 0.0;
    for (size_t i = 0; i < groupNum; i++) {
      float perGroupSum = 0.0;
      for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if ((c.getType() == "CPU" || c.getType() == "GPU" ||
             c.getType() == "CGRA") &&
            c.getGroup() == i) {
          perGroupSum += c.getDutyRatio();
        }
      }

      Value selCPU = rewriter.create<arith::ConstantIndexOp>(loc, i);
      Value isCPU = rewriter.create<arith::CmpIOp>(
          loc, arith::CmpIPredicate::eq, selCPU, cpuId);
      rewriter.create<scf::IfOp>(
          loc, isCPU, [&](OpBuilder &builder, Location loc) {
            Value lowerRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());
            sum += perGroupSum;
            Value upperRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());

            Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
            Value lowerI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), lowerF);
            Value lowerBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), lowerI);

            Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
            Value upperI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), upperF);
            Value upperBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), upperI);
            if (i == groupNum - 1) {
              upperBound = numIdx;
            }

            Value groupNumIdx = rewriter.create<arith::SubIOp>(
                loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);
            Value groupSizeIn = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/msgLenIdx);
            auto byte_shift =
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/msgLenIdx);
            Value groupMemIn = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8Type, /*source=*/candidates,
                /*byte_shift=*/
                byte_shift,
                /*sizes=*/ValueRange{groupSizeIn});
            auto ENV_printschedule = std::getenv("PRINT_SCHEDULE");
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemIn, msgLen, 
                                    static_cast<int>(OP_TYPE::INPUT_CRACK_SM3));

            // 'hyper.for' to run multiple no-dependency MD5 tasks
            auto hyperForOp = rewriter.create<hyper::ForOp>(
                loc, /*lowerBound=*/c0, /*upperBound=*/groupNumIdx,
                /*memIns=*/ValueRange{groupMemIn, target},
                /*memOuts=*/ValueRange{initMemref},
                /*odsArg4=*/
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange miargs, ValueRange moargs) {
                  Value idx = builder.create<arith::MulIOp>(loc, /*lhs=*/iv,
                                                            /*rhs=*/msgLenIdx);
                  // allocate temporary memory for hash result
                  Value alloca = builder.create<memref::AllocaOp>(
                      loc, /*memrefType=*/I8x32Type,
                      /*dynamicSizes=*/ValueRange{});
                  // get original candidate message
                  Value oriMsg = builder.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8Type, /*source=*/miargs[0],
                      /*byte_shift=*/idx, /*sizes=*/ValueRange{msgLenIdx});
                  // get memref to be ready for calculation of MD5
                  Value partMemIn = builder.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8Type, /*source=*/miargs[0],
                      /*byte_shift=*/idx, /*sizes=*/ValueRange{msgLenIdx});

                  // MD5 computation
                  rewriter.create<crypto::SM3OP>(loc, /*input=*/partMemIn,
                                                 /*msgLen=*/msgLen,
                                                 /*output=*/alloca);

                  // compare with target
                  Value cond = builder.create<crypto::CmpOp>(
                      loc, /*isEqual*/ builder.getI1Type(), /*src=*/alloca,
                      /*dst=*/miargs[1]);
                  builder.create<scf::IfOp>(
                      loc, /*cond=*/cond, /*thenBuilder=*/
                      [&](OpBuilder &builder, Location loc) {
                        // do nothing
                        builder.create<scf::YieldOp>(loc);
                      },
                      /*elseBuilder=*/
                      [&](OpBuilder &builder, Location loc) {
                        // store the result
                        builder.create<scf::ForOp>(
                            loc, /*lowerBound=*/c0, /*upperBound=*/msgLenIdx,
                            /*step=*/c1, /*iterArgs=*/ValueRange{},
                            /*bodyBuilder=*/
                            [&](OpBuilder &builder, Location loc, Value iv,
                                ValueRange args) {
                              // store the result
                              builder.create<memref::StoreOp>(
                                  loc, /*value=*/c0i8, /*memref=*/oriMsg,
                                  /*indices=*/ValueRange{iv});
                              builder.create<scf::YieldOp>(loc);
                            });
                        builder.create<scf::YieldOp>(loc);
                      });

                  // reduce the results
                  builder.create<hyper::ReduceOp>(
                      loc, /*operands=*/ValueRange{oriMsg},
                      /*bodyBuilderFn=*/
                      ArrayRef<mlir::function_ref<void(OpBuilder &, Location,
                                                       Value, Value)>>{
                          [&](OpBuilder &builder, Location loc, Value lhs,
                              Value rhs) {
                            builder.create<scf::ForOp>(
                                loc, /*lowerBound=*/c0,
                                /*upperBound=*/msgLenIdx, /*step=*/c1,
                                /*iterArgs=*/ValueRange{},
                                /*bodyBuilderFns=*/
                                [&](OpBuilder &builder, Location loc, Value iv,
                                    ValueRange args) {
                                  // load one byte
                                  Value val = builder.create<memref::LoadOp>(
                                      loc, /*memref=*/rhs,
                                      /*indices=*/ValueRange{iv});
                                  // atomic "or" operation
                                  builder.create<memref::AtomicRMWOp>(
                                      loc, /*kind=*/
                                      arith::AtomicRMWKindAttr::get(
                                          op.getContext(),
                                          arith::AtomicRMWKind::ori),
                                      /*value=*/val, /*memref=*/lhs,
                                      /*indices=*/ValueRange{iv});
                                  builder.create<scf::YieldOp>(loc);
                                });
                            builder.create<hyper::ReduceReturnOp>(
                                loc, /*result=*/lhs);
                          }});

                  builder.create<hyper::YieldOp>(loc);
                });

            if (ENV_printschedule && std::string(ENV_printschedule) == "ON") {
              mlirPrintScheduleInfo(rewriter, op, i, target, msgLen, 
                                    static_cast<int>(OP_TYPE::EXTRA_CRACK_SM3));
              mlirPrintScheduleInfo(rewriter, op, i, initMemref, msgLen, 
                                    static_cast<int>(OP_TYPE::OUTPUT_CRACK_SM3));
            }

            SmallVector<Attribute> deviceSplit;
            for (wafer::Chip c :
                 HardwareInfo.getArchitecture().getChips().getChipArray()) {
              float tmpRatio;
              if ((c.getType() == "CPU" || c.getType() == "GPU" || c.getType() == "CGRA") &&
                  c.getGroup() == i) {
                assert(perGroupSum != 0.0 && "exist lazy group!");
                tmpRatio = c.getDutyRatio() / perGroupSum;
                if (c.getType() == "CPU") {
                  // cpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else if (c.getType() == "CGRA") {
                  // cpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cgra")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else {
                  // gpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("gpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                }
              }
            } // end of chip split for-loop

            // add additional 'hyper.devices' attribute for 'hyper.for' op
            auto deviceSplitAttr = rewriter.getNamedAttr(
                "devices", rewriter.getArrayAttr(deviceSplit));
            // customized 'isSharedMem' attribute for 'hyper.for' op
            auto isSharedMemAttr = rewriter.getNamedAttr(
                "isSharedMem", rewriter.getBoolArrayAttr({false, true, true}));

            SmallVector<NamedAttribute> attrs;
            for (auto attr : hyperForOp->getAttrs()) {
              if (attr.getName() != "devices" &&
                  attr.getName() != "isSharedMem")
                attrs.push_back(attr);
            }
            attrs.push_back(deviceSplitAttr);
            attrs.push_back(isSharedMemAttr);
            hyperForOp->setAttrs(attrs);

            rewriter.create<scf::YieldOp>(loc);
          });
    }

    // Determine if a cpu is a master cpu(0)
    Value isMasterCpu = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/cpuId, /*rhs=*/c0);

    // Determine if crack result is null
    Value Result0 = rewriter.create<memref::LoadOp>(loc, /*memref=*/initMemref,
                                                    /*indices=*/ValueRange{c0});
    Value isNullResult = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/Result0,
        /*rhs=*/c0i8);

    // when cpu is a master cpu and crack result is null, cpu get crack result
    Value isMasterAndNull =
        rewriter.create<arith::AndIOp>(loc, isMasterCpu, isNullResult);
    rewriter.create<scf::IfOp>(
        loc, /*cond=*/isMasterAndNull,
        /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
          for (size_t i = 1; i < groupNum; i++) {
            Value srcid = rewriter.create<arith::ConstantIndexOp>(loc, i);
            rewriter.create<sst::MsgRecvOp>(loc, /*src=*/srcid, /*dst=*/c0,
                                            /*buffer=*/initMemref,
                                            /*size=*/msgLenIdx);
          }
          rewriter.create<scf::YieldOp>(loc);
        });

    // when cpu is not a master cpu and crack result is not null, cpu send crack
    // result Negate isMasterCpu and isNullResult
    Value notMasterCpu = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::ne, /*lhs=*/cpuId, /*rhs=*/c0);
    Value notNullResult = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::ne, /*lhs=*/Result0,
        /*rhs=*/c0i8);
    Value notMasterAndnotNull =
        rewriter.create<arith::AndIOp>(loc, notMasterCpu, notNullResult);
    rewriter.create<scf::IfOp>(
        loc, /*cond=*/notMasterAndnotNull,
        /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
          rewriter.create<sst::MsgSendOp>(loc, /*src=*/cpuId, /*dst=*/c0,
                                          /*buffer=*/initMemref,
                                          /*size=*/msgLenIdx);
          rewriter.create<scf::YieldOp>(loc);
        });

    // replace the original operation with the result
    rewriter.replaceOp(op, ValueRange{initMemref});
    // rewriter.eraseOp(op);
    return success();
  }

}; // end of MD5Crack

} // namespace

//===----------------------------------------------------------------------===//
// AES
//===----------------------------------------------------------------------===//

namespace {
// Little-endian
SmallVector<int> sBoxAndMixColumnMatrix = {
    2071428195,
    /*0x7b777c63*/ -982553614,
    /*0xc56f6bf2*/ 728170800,
    /*0x2b670130*/ 1990973438, /*0x76abd7fe*/
    2110358218,
    /*0x7dc982ca*/ -263759366,
    /*0xf04759fa*/ -1348283219,
    /*0xafa2d4ad*/ -1066228580, /*0xc072a49c*/
    647232951,
    /*0x2693fdb7*/ -856211658,
    /*0xccf73f36*/ -236608204,
    /*0xf1e5a534*/ 355588209, /*0x1531d871*/
    -1021065468,
    /*0xc323c704*/ -1710909928,
    /*0x9a059618*/ -494923257,
    /*0xe2801207*/ 1974609899, /*0x75b227eb*/
    439124745,
    /*0x1a2c8309*/ -1604686309,
    /*0xa05a6e1b*/ -1277805742,
    /*0xb3d63b52*/ -2077236439, /*0x842fe329*/
    -318713517,
    /*0xed00d153*/ 1538391072,
    /*0x5bb1fc20*/ 968805226,
    /*0x39becb6a*/ -816296886, /*0xcf584c4a*/
    -72683568,
    /*0xfbaaefd0*/ -2060235453,
    /*0x85334d43*/ 2130901317,
    /*0x7f02f945*/ -1465959344, /*0xa89f3c50*/
    -1891589295,
    /*0x8f40a351*/ -180839022,
    /*0xf5389d92*/ 567981756,
    /*0x21dab6bc*/ -755761392, /*0xd2f3ff10*/
    -334295859,
    /*0xec130ccd*/ 390371167,
    /*0x1744975f*/ 1031710660,
    /*0x3d7ea7c4*/ 1931042148, /*0x73195d64*/
    -598769312,
    /*0xdc4f8160*/ -2003817950,
    /*0x88902a22*/ 347663942,
    /*0x14b8ee46*/ -620011810, /*0xdb0b5ede*/
    171586272,
    /*0x0a3a32e0*/ 1545864777,
    /*0x5c240649*/ 1655493570,
    /*0x62acd3c2*/ 2045023633, /*0x79e49591*/
    1832372455,
    /*0x6d37c8e7*/ -1454451315,
    /*0xa94ed58d*/ -353085844,
    /*0xeaf4566c*/ 145652325, /*0x08ae7a65*/
    774207674,
    /*0x2e2578ba*/ -961239524,
    /*0xc6b4a61c*/ 527752680,
    /*0x1f74dde8*/ -1970553525, /*0x8a8bbd4b*/
    1723154032,
    /*0x66b53e70*/ 251003720,
    /*0x0ef60348*/ -1185467039,
    /*0xb9573561*/ -1642217082, /*0x9e1dc186*/
    295237857,
    /*0x1198f8e1*/ -1802577559,
    /*0x948ed969*/ -377020773,
    /*0xe9871e9b*/ -551004722, /*0xdf2855ce*/
    227123596,
    /*0x0d89a18c*/ 1749214911,
    /*0x6842e6bf*/ 254646593,
    /*0x0f2d9941*/ 381375664, /*0x16bb54b0*/
    16843522,
    /*0x01010302*/ 16974337,
    /*0x01030201*/ 50462977,
    /*0x03020101*/ 33620227, /*0x02010103*/
};

SmallVector<int> invSBoxAndMixColumnMatrix = {
    -714471086,
    /*0xd56a0952*/ 950351408,
    /*0x38a53630*/ -1633468225,
    /*0x9ea340bf*/ -69733503, /*0xfbd7f381*/
    -2110135428,
    /*0x8239e37c*/ -2013319269,
    /*0x87ff2f9b*/ 1145278004,
    /*0x44438e34*/ -873865532, /*0xcbe9dec4*/
    848591700,
    /*0x32947b54*/ 1025753766,
    /*0x3d23c2a6*/ 194333934,
    /*0x0b954cee*/ 1321466434, /*0x4ec3fa42*/
    1721839112,
    /*0x66a12e08*/ -1306207960,
    /*0xb224d928*/ 1235377014,
    /*0x49a25b76*/ 634489709, /*0x25d18b6d*/
    1693907058,
    /*0x64f6f872*/ 379086982,
    /*0x16986886*/ -866343724,
    /*0xcc5ca4d4*/ -1833540259, /*0x92b6655d*/
    1346924652,
    /*0x5048706c*/ -625349123,
    /*0xdab9edfd*/ 1464210782,
    /*0x5746155e*/ -2070049369, /*0x849d8da7*/
    11262096,
    /*0x00abd890*/ 181648524,
    /*0x0ad3bc8c*/ 89711863,
    /*0x0558e4f7*/ 105231288, /*0x0645b3b8*/
    -1893847856,
    /*0x8f1e2cd0*/ 34553802,
    /*0x020f3fca*/ 62762945,
    /*0x03bdafc1*/ 1804210945, /*0x6b8a1301*/
    1091670330,
    /*0x4111913a*/ -354654385,
    /*0xeadc674f*/ -825232745,
    /*0xcecff297*/ 1944499440, /*0x73e6b4f0*/
    578071702,
    /*0x2274ac96*/ -2060079641,
    /*0x8535ade7*/ -398984734,
    /*0xe837f9e2*/ 1860138268, /*0x6edf751c*/
    1897591111,
    /*0x711af147*/ -1983567587,
    /*0x89c5291d*/ 241350511,
    /*0x0e62b76f*/ 465442986, /*0x1bbe18aa*/
    1262376700,
    /*0x4b3e56fc*/ 544854726,
    /*0x2079d2c6*/ -20915302,
    /*0xfec0db9a*/ -195375752, /*0xf45acd78*/
    866704671,
    /*0x33a8dd1f*/ 835127176,
    /*0x31c70788*/ 1494225585,
    /*0x591012b1*/ 1609334823, /*0x5fec8027*/
    -1451273888,
    /*0xa97f5160*/ 222999833,
    /*0x0d4ab519*/ -1619335891,
    /*0x9f7ae52d*/ -274937453, /*0xef9cc993*/
    1295769760,
    /*0x4d3be0a0*/ -1326110034,
    /*0xb0f52aae*/ 1018948552,
    /*0x3cbbebc8*/ 1637438339, /*0x61995383*/
    2114202391,
    /*0x7e042b17*/ 651589562,
    /*0x26d677ba*/ 1662282209,
    /*0x631469e1*/ 2097946965, /*0x7d0c2155*/
    151849742,
    /*0x090d0b0e*/ 218828297,
    /*0x0d0b0e09*/ 185469197,
    /*0x0b0e090d*/ 235474187, /*0x0e090d0b*/
};

void keyExpansion(Location loc, OpBuilder &rewriter, Value keyBase,
                  Value sBoxBase) {
  SmallVector<int64_t, 4> inv4{1, 2, 3, 0};
  Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
  Value c4 = rewriter.create<arith::ConstantIndexOp>(loc, 4);
  VectorType vectorI1x4 = VectorType::get({4}, rewriter.getI1Type());
  VectorType vectorI8x4 = VectorType::get({4}, rewriter.getI8Type());
  VectorType vectorI32x4 = VectorType::get({4}, rewriter.getI32Type());
  Value maskI1x4 = rewriter.create<vector::CreateMaskOp>(loc, vectorI1x4, c4);
  unsigned char rcon[11] = {0x00, 0x01, 0x02, 0x04, 0x08, 0x10,
                            0x20, 0x40, 0x80, 0x1b, 0x36};

  for (size_t i = 16; i < 176; i += 4) {
    Value keyLast1 = rewriter.create<vector::LoadOp>(
        loc, vectorI8x4, keyBase,
        ValueRange{rewriter.create<arith::ConstantIndexOp>(loc, i - 4)});
    if (!(i % 16)) {
      keyLast1 =
          rewriter.create<vector::ShuffleOp>(loc, keyLast1, keyLast1, inv4);

      Value keyLast1Index =
          rewriter.create<arith::ExtUIOp>(loc, vectorI32x4, keyLast1);
      keyLast1 = rewriter.create<vector::GatherOp>(
          loc, vectorI8x4, sBoxBase, c0, keyLast1Index, maskI1x4, keyLast1);

      Value rconValue = rewriter.create<arith::ConstantOp>(
          loc, rewriter.getI8Type(), rewriter.getI8IntegerAttr(rcon[i / 16]));
      Value keyValue =
          rewriter.create<vector::ExtractElementOp>(loc, keyLast1, c0);
      keyLast1 = rewriter.create<vector::InsertElementOp>(
          loc, vectorI8x4,
          rewriter.create<arith::XOrIOp>(loc, keyValue, rconValue), keyLast1,
          c0);
    }
    Value keyLast4 = rewriter.create<vector::LoadOp>(
        loc, vectorI8x4, keyBase,
        ValueRange{rewriter.create<arith::ConstantIndexOp>(loc, i - 16)});
    Value result =
        rewriter.create<arith::XOrIOp>(loc, vectorI8x4, keyLast1, keyLast4);
    rewriter.create<vector::StoreOp>(
        loc, result, keyBase,
        ValueRange{rewriter.create<arith::ConstantIndexOp>(loc, i)});
  }
}

class CryptoAESToHyperPattern : public OpRewritePattern<crypto::AESmultiOP> {
public:
  using OpRewritePattern<crypto::AESmultiOP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::AESmultiOP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    Value allMsgs = op->getOperand(0);      // memref<?xi8>
    Value msgLen = op->getOperand(1);       // i32
    Value outputBuffer = op->getOperand(2); // memref<?xi8>
    Value cptLen = op->getOperand(3);       // i32
    Value allkeys = op->getOperand(4);      // memref<?xi8>
    Value keyLen = op->getOperand(5);       // i32
    Value msgCounts = op->getOperand(6);    // i32

    // types
    auto dynI8MemrefType =
        MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    // constants
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/1);
    Value c16 = rewriter.create<arith::ConstantIndexOp>(loc, 16);
    // index cast
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgLen);
    Value cptLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/cptLen);
    Value keyLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/keyLen);
    Value numIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgCounts);
    Value Nfp =
        rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(), numIdx);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);

    // first we need cpuid info
    Value cpuId = rewriter.create<sst::GetCpuIdOp>(
        loc, /*cpuId=*/rewriter.getIndexType());

    wafer::Root &HardwareInfo = wafer::Root::getInstance();

    std::size_t groupNum = HardwareInfo.getArchitecture().getChips().CPUNum();

    // if sst use gpu, we use 'RegisterFatbin' and 'RegisterFunc' op.
    bool useGPU = false;
    bool useCGRA = false;
    for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
      if (c.getType() == "GPU" && c.getDutyRatio() > 0) {
        useGPU = true;
      }else if (c.getType() == "CGRA" && c.getDutyRatio() > 0) {
        useCGRA = true;
      }
    }
    if (useGPU) {
      Value handle = rewriter.create<sst::RegisterFatbinOp>(
          loc, /*type=*/rewriter.getIndexType());
      rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle,
                                           /*hostFunc=*/c1);
    }
    if (useCGRA) {
      convert();
    }

    float sum = 0.0;
    for (size_t i = 0; i < groupNum; i++) {
      SmallVector<Attribute> deviceSplit;
      float perGroupSum = 0.0;
      for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if ((c.getType() == "CPU" || c.getType() == "GPU" ||
             c.getType() == "CGRA") &&
            c.getGroup() == i) {
          perGroupSum += c.getDutyRatio();
        }
      }

      Value SBoxAndMixColumnMatrix = rewriter.create<arith::ConstantOp>(
          loc, rewriter.getI32VectorAttr(sBoxAndMixColumnMatrix));
      SBoxAndMixColumnMatrix = rewriter.create<vector::BitCastOp>(
          loc, VectorType::get({272}, rewriter.getI8Type()),
          SBoxAndMixColumnMatrix);
      Value memSize = rewriter.create<arith::ConstantIndexOp>(loc, 272);
      Value sBoxBase =
          rewriter.create<memref::AllocOp>(loc, dynI8MemrefType, memSize);
      rewriter.create<vector::StoreOp>(loc, SBoxAndMixColumnMatrix, sBoxBase,
                                       c0);

      Value keySize = rewriter.create<arith::ConstantIndexOp>(loc, 176);
      Value keyBase =
          rewriter.create<memref::AllocOp>(loc, dynI8MemrefType, keySize);
      VectorType vectorI8x16 = VectorType::get({16}, rewriter.getI8Type());
      Value plainKey =
          rewriter.create<vector::LoadOp>(loc, vectorI8x16, allkeys, c0);
      rewriter.create<vector::StoreOp>(loc, plainKey, keyBase, c0);
      keyExpansion(loc, rewriter, keyBase, sBoxBase);

      Value selCPU = rewriter.create<arith::ConstantIndexOp>(loc, i);
      Value isCPU = rewriter.create<arith::CmpIOp>(
          loc, arith::CmpIPredicate::eq, selCPU, cpuId);
      rewriter.create<scf::IfOp>(
          loc, isCPU, [&](OpBuilder &builder, Location loc) {
            Value lowerRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());
            sum += perGroupSum;
            Value upperRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());

            Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
            Value lowerI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), lowerF);
            Value lowerBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), lowerI);

            Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
            Value upperI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), upperF);
            Value upperBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), upperI);
            if (i == groupNum - 1) {
              upperBound = numIdx;
            }
            Value groupNumIdx = rewriter.create<arith::SubIOp>(
                loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);
            Value groupSizeIn = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/msgLenIdx);
            Value groupSizeOut = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/cptLenIdx);
            Value byte_shift =
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/msgLenIdx);
            Value groupMemIn = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/allMsgs,
                /*byte_shift=*/
                byte_shift,
                /*sizes=*/ValueRange{groupSizeIn});
            Value groupMemOut = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/outputBuffer,
                /*byte_shift=*/
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/cptLenIdx),
                /*sizes=*/ValueRange{groupSizeOut});
            Value key = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/allkeys,
                /*byte_shift=*/c0, /*sizes=*/ValueRange{keyLenIdx});
            auto ENV_printschedule = std::getenv("PRINT_SCHEDULE");
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemIn, msgLen, 
                                    static_cast<int>(OP_TYPE::INPUT_MULTI_AES));

            // 'hyper.for' to run multiple no-dependency AES tasks
            auto hyperForOp = rewriter.create<hyper::ForOp>(
                loc, /*lowerBound=*/c0, /*upperBound=*/groupNumIdx,
                /*memIns=*/ValueRange{groupMemIn},
                /*memOuts=*/ValueRange{groupMemOut},
                /*odsArg4=*/
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange miargs, ValueRange moargs) {
                  Value inShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/msgLenIdx);
                  Value outShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/cptLenIdx);
                  // extract a single message out
                  Value input = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/miargs[0], /*byte_shift=*/inShift,
                      /*sizes=*/ValueRange{msgLenIdx});
                  Value output = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/moargs[0], /*byte_shift=*/outShift,
                      /*sizes=*/ValueRange{cptLenIdx});
                  // meaningless code, just for sst can get keyBase and sBoxBase
                  // which is out of the "hyper.for"
                  Value meaningless =
                      rewriter.create<memref::RankOp>(loc, /*memref=*/keyBase);
                  meaningless =
                      rewriter.create<memref::RankOp>(loc, /*memref=*/sBoxBase);
                  // AES computation
                  // if (keyLenIdx == c16){
                  rewriter.create<crypto::AESOP>(
                      loc, /*input=*/input, /*msgLen=*/msgLen,
                      /*output=*/output, /*key=*/keyBase,
                      /*sBoxBase=*/sBoxBase);
                  // } else if (keyLenIdx == c24) {
                  // } else if (keyLenIdx == c32) { }
                  builder.create<hyper::YieldOp>(loc);
                });

            if (ENV_printschedule && std::string(ENV_printschedule) == "ON") {
              mlirPrintScheduleInfo(rewriter, op, i, allkeys, cptLen, 
                                    static_cast<int>(OP_TYPE::EXTRA_MULTI_CRACK_AES));
              mlirPrintScheduleInfo(rewriter, op, i, groupMemOut, cptLen, 
                                    static_cast<int>(OP_TYPE::OUTPUT_MULTI_AES));
            }


            for (wafer::Chip c :
                 HardwareInfo.getArchitecture().getChips().getChipArray()) {
              float tmpRatio;
              if ((c.getType() == "CPU" || c.getType() == "GPU" || c.getType() == "CGRA") &&
                  c.getGroup() == i) {
                assert(perGroupSum != 0.0 && "exist lazy group!");
                tmpRatio = c.getDutyRatio() / perGroupSum;
                if (c.getType() == "CPU") {
                  // cpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else if (c.getType() == "CGRA") {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cgra")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else {
                  // gpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("gpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                }
              }
            } // end of chip split for-loop

            // add additional 'hyper.devices' attribute for 'hyper.for' op
            auto deviceSplitAttr = rewriter.getNamedAttr(
                "devices", rewriter.getArrayAttr(deviceSplit));
            SmallVector<NamedAttribute> attrs;
            for (auto attr : hyperForOp->getAttrs()) {
              if (attr.getName() !=
                  "devices") // overwrite default 'hyper.devices' attribute
                attrs.push_back(attr);
            }
            attrs.push_back(deviceSplitAttr);
            hyperForOp->setAttrs(attrs);

            rewriter.create<scf::YieldOp>(loc);
          });
    }

    // erase the original 'crypto.aesmulti' op
    rewriter.eraseOp(op);

    return success();
  }
};

class AESCrackLoweringPattern : public OpRewritePattern<crypto::AESCrackOp> {
public:
  using OpRewritePattern<crypto::AESCrackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::AESCrackOp op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    Value allMsgs = op->getOperand(0);      // memref<?xi8>
    Value msgLen = op->getOperand(1);       // i32
    Value outputBuffer = op->getOperand(2); // memref<?xi8>
    Value cptLen = op->getOperand(3);       // i32
    Value allkeys = op->getOperand(4);      // memref<?xi8>
    Value keyLen = op->getOperand(5);       // i32
    Value msgCounts = op->getOperand(6);    // i32

    // types
    auto dynI8MemrefType =
        MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    // constants
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/1);
    Value c16 = rewriter.create<arith::ConstantIndexOp>(loc, 16);
    // index cast
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgLen);
    Value cptLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/cptLen);
    Value keyLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/keyLen);
    Value numIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgCounts);
    Value Nfp =
        rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(), numIdx);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);

    // first we need cpuid info
    Value cpuId = rewriter.create<sst::GetCpuIdOp>(
        loc, /*cpuId=*/rewriter.getIndexType());

    wafer::Root &HardwareInfo = wafer::Root::getInstance();

    std::size_t groupNum = HardwareInfo.getArchitecture().getChips().CPUNum();

    // if sst use gpu, we use 'RegisterFatbin' and 'RegisterFunc' op.
    bool useGPU = false;
    bool useCGRA = false;
    for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
      if (c.getType() == "GPU" && c.getDutyRatio() > 0) {
        useGPU = true;
      }else if (c.getType() == "CGRA" && c.getDutyRatio() > 0) {
        useCGRA = true;
      }
    }
    if (useGPU) {
      Value handle = rewriter.create<sst::RegisterFatbinOp>(
          loc, /*type=*/rewriter.getIndexType());
      rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle,
                                           /*hostFunc=*/c1);
    }
    if (useCGRA) {
      convert();
    }

    float sum = 0.0;
    for (size_t i = 0; i < groupNum; i++) {
      float perGroupSum = 0.0;
      for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if ((c.getType() == "CPU" || c.getType() == "GPU" ||
             c.getType() == "CGRA") &&
            c.getGroup() == i) {
          perGroupSum += c.getDutyRatio();
        }
      }

      Value SBoxAndMixColumnMatrix = rewriter.create<arith::ConstantOp>(
          loc, rewriter.getI32VectorAttr(sBoxAndMixColumnMatrix));
      SBoxAndMixColumnMatrix = rewriter.create<vector::BitCastOp>(
          loc, VectorType::get({272}, rewriter.getI8Type()),
          SBoxAndMixColumnMatrix);
      Value memSize = rewriter.create<arith::ConstantIndexOp>(loc, 272);
      Value sBoxBase =
          rewriter.create<memref::AllocOp>(loc, dynI8MemrefType, memSize);
      rewriter.create<vector::StoreOp>(loc, SBoxAndMixColumnMatrix, sBoxBase,
                                       c0);

      Value keySize = rewriter.create<arith::ConstantIndexOp>(loc, 176);
      Value keyBase =
          rewriter.create<memref::AllocOp>(loc, dynI8MemrefType, keySize);
      VectorType vectorI8x16 = VectorType::get({16}, rewriter.getI8Type());
      Value plainKey =
          rewriter.create<vector::LoadOp>(loc, vectorI8x16, allkeys, c0);
      rewriter.create<vector::StoreOp>(loc, plainKey, keyBase, c0);
      keyExpansion(loc, rewriter, keyBase, sBoxBase);

      Value InvSBoxAndMixColumnMatrix = rewriter.create<arith::ConstantOp>(
          loc, rewriter.getI32VectorAttr(invSBoxAndMixColumnMatrix));
      InvSBoxAndMixColumnMatrix = rewriter.create<vector::BitCastOp>(
          loc, VectorType::get({272}, rewriter.getI8Type()),
          InvSBoxAndMixColumnMatrix);
      rewriter.create<vector::StoreOp>(loc, InvSBoxAndMixColumnMatrix, sBoxBase,
                                       c0);

      Value selCPU = rewriter.create<arith::ConstantIndexOp>(loc, i);
      Value isCPU = rewriter.create<arith::CmpIOp>(
          loc, arith::CmpIPredicate::eq, selCPU, cpuId);
      rewriter.create<scf::IfOp>(
          loc, isCPU, [&](OpBuilder &builder, Location loc) {
            Value lowerRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());
            sum += perGroupSum;
            Value upperRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());

            Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
            Value lowerI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), lowerF);
            Value lowerBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), lowerI);

            Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
            Value upperI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), upperF);
            Value upperBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), upperI);
            if (i == groupNum - 1) {
              upperBound = numIdx;
            }

            Value groupNumIdx = rewriter.create<arith::SubIOp>(
                loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);
            Value groupSizeIn = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/msgLenIdx);
            Value groupSizeOut = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/cptLenIdx);
            Value byte_shift =
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/msgLenIdx);
            Value groupMemIn = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/allMsgs,
                /*byte_shift=*/
                byte_shift,
                /*sizes=*/ValueRange{groupSizeIn});
            Value groupMemOut = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/outputBuffer,
                /*byte_shift=*/
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/cptLenIdx),
                /*sizes=*/ValueRange{groupSizeOut});
            Value key = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/allkeys,
                /*byte_shift=*/c0, /*sizes=*/ValueRange{keyLenIdx});
            auto ENV_printschedule = std::getenv("PRINT_SCHEDULE");
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemIn, msgLen, 
                                    static_cast<int>(OP_TYPE::INPUT_CRACK_AES));

            // 'hyper.for' to run multiple no-dependency AES tasks
            auto hyperForOp = rewriter.create<hyper::ForOp>(
                loc, /*lowerBound=*/c0, /*upperBound=*/groupNumIdx,
                /*memIns=*/ValueRange{groupMemIn},
                /*memOuts=*/ValueRange{groupMemOut},
                /*odsArg4=*/
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange miargs, ValueRange moargs) {
                  Value inShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/msgLenIdx);
                  Value outShift = rewriter.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/cptLenIdx);
                  // extract a single message out
                  Value input = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/miargs[0], /*byte_shift=*/inShift,
                      /*sizes=*/ValueRange{msgLenIdx});
                  Value output = rewriter.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/moargs[0], /*byte_shift=*/outShift,
                      /*sizes=*/ValueRange{cptLenIdx});
                  // meaningless code, just for sst can get keyBase and sBoxBase
                  // which is out of the "hyper.for"
                  Value meaningless =
                      rewriter.create<memref::RankOp>(loc, /*memref=*/keyBase);
                  meaningless =
                      rewriter.create<memref::RankOp>(loc, /*memref=*/sBoxBase);
                  // AES computation
                  // if (keyLenIdx == c16){
                  rewriter.create<crypto::AESOP>(
                      loc, /*input=*/input, /*msgLen=*/msgLen,
                      /*output=*/output, /*key=*/keyBase,
                      /*sBoxBase=*/sBoxBase);
                  // } else if (keyLenIdx == c24) {
                  // } else if (keyLenIdx == c32) { }
                  builder.create<hyper::YieldOp>(loc);
                });

            if (ENV_printschedule && std::string(ENV_printschedule) == "ON") {
              mlirPrintScheduleInfo(rewriter, op, i, allkeys, cptLen, 
                                    static_cast<int>(OP_TYPE::EXTRA_MULTI_CRACK_AES));
              mlirPrintScheduleInfo(rewriter, op, i, groupMemOut, cptLen, 
                                    static_cast<int>(OP_TYPE::OUTPUT_CRACK_AES));
            }

            SmallVector<Attribute> deviceSplit;
            for (wafer::Chip c :
                 HardwareInfo.getArchitecture().getChips().getChipArray()) {
              float tmpRatio;
              if ((c.getType() == "CPU" || c.getType() == "GPU" || c.getType() == "CGRA") &&
                  c.getGroup() == i) {
                assert(perGroupSum != 0.0 && "exist lazy group!");
                tmpRatio = c.getDutyRatio() / perGroupSum;
                if (c.getType() == "CPU") {
                  // cpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                }  else if (c.getType() == "CGRA") {
                  // cpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cgra")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else {
                  // gpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("gpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                }
              }
            } // end of chip split for-loop

            // add additional 'hyper.devices' attribute for 'hyper.for' op
            auto deviceSplitAttr = rewriter.getNamedAttr(
                "devices", rewriter.getArrayAttr(deviceSplit));
            SmallVector<NamedAttribute> attrs;
            for (auto attr : hyperForOp->getAttrs()) {
              if (attr.getName() !=
                  "devices") // overwrite default 'hyper.devices' attribute
                attrs.push_back(attr);
            }
            attrs.push_back(deviceSplitAttr);
            hyperForOp->setAttrs(attrs);

            rewriter.create<scf::YieldOp>(loc);
          });
    }

    // erase the original 'crypto.aes.crack' op
    rewriter.eraseOp(op);

    return success();
  }

}; // end of AESCrack


// 使用aes加密算法进行破译
class AESCrackLoweringPattern_ : public OpRewritePattern<crypto::AESCrackOp> {
public:
  using OpRewritePattern<crypto::AESCrackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::AESCrackOp op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    Value candidates = op->getOperand(0);      // memref<?xi8>
    Value msgLen = op->getOperand(1);       // i32
    Value targetOutput = op->getOperand(2); // memref<?xi8>
    Value cptLen = op->getOperand(3);       // i32
    Value candidatesKeys = op->getOperand(4);      // memref<?xi8>
    Value keyLen = op->getOperand(5);       // i32
    Value msgCounts = op->getOperand(6);    // i32
    // types
    auto dynI8MemrefType =
        MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    auto I8x16MemrefType = MemRefType::get(ArrayRef<int64_t>{16}, rewriter.getI8Type());
    // constants
    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, /*value=*/1);
    Value c16 = rewriter.create<arith::ConstantIndexOp>(loc, 16);
    Value c0i8 = rewriter.create<arith::ConstantOp>(loc, rewriter.getI8IntegerAttr(0));
    auto candidates_key_len = c16;
    // index cast
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgLen);
    Value cptLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/cptLen);
    Value keyLenIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/keyLen);
    Value numIdx = rewriter.create<arith::IndexCastOp>(
        loc, /*out=*/rewriter.getIndexType(), /*in=*/msgCounts);
    Value Nfp =
        rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(), numIdx);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);

    // count candidate keys
    auto candidate_keys_total_len = rewriter.create<memref::DimOp>(loc, candidatesKeys, 0);
    auto index_16 = rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(16));
    auto candidate_keys_count = rewriter.create<index::DivUOp>(loc, candidate_keys_total_len, index_16);

    // memref with initial values to store the crack result
    auto allocOp = rewriter.create<memref::AllocOp>(
        loc, /*memrefType=*/dynI8MemrefType, /*dynamicSizes=*/ValueRange{msgLenIdx});
    allocOp->setAttr("operand_segment_sizes",
                     rewriter.getDenseI32ArrayAttr({1, 0}));
    Value initMemref = allocOp.getResult();
    rewriter.create<scf::ForOp>(
        loc, /*lowerBound=*/c0, /*upperBound=*/msgLenIdx, /*step=*/c1,
        /*iterArgs=*/ValueRange{}, /*bodyBuilder=*/
        [&](OpBuilder &builder, Location loc, Value ivs, ValueRange args) {
          builder.create<memref::StoreOp>(loc, /*valueToStore=*/c0i8,
                                          /*memref=*/initMemref,
                                          /*indices=*/ValueRange{ivs});
          builder.create<scf::YieldOp>(loc);
        });


    // first we need cpuid info
    Value cpuId = rewriter.create<sst::GetCpuIdOp>(
        loc, /*cpuId=*/rewriter.getIndexType());

    wafer::Root &HardwareInfo = wafer::Root::getInstance();

    std::size_t groupNum = HardwareInfo.getArchitecture().getChips().CPUNum();

    // if sst use gpu, we use 'RegisterFatbin' and 'RegisterFunc' op.
    bool useGPU = false;
    bool useCGRA = false;
    for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
      if (c.getType() == "GPU" && c.getDutyRatio() > 0) {
        useGPU = true;
      }else if (c.getType() == "CGRA" && c.getDutyRatio() > 0) {
        useCGRA = true;
      }
    }
    if (useGPU) {
      Value handle = rewriter.create<sst::RegisterFatbinOp>(
          loc, /*type=*/rewriter.getIndexType());
      rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle,
                                           /*hostFunc=*/c1);
    }
    if (useCGRA) {
      convert();
    }

    float sum = 0.0;
    for (size_t i = 0; i < groupNum; i++) {
      SmallVector<Attribute> deviceSplit;
      float perGroupSum = 0.0;
      for (auto c : HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if ((c.getType() == "CPU" || c.getType() == "GPU" ||
             c.getType() == "CGRA") &&
            c.getGroup() == i) {
          perGroupSum += c.getDutyRatio();
        }
      }

      // sbox
      Value SBoxAndMixColumnMatrix = rewriter.create<arith::ConstantOp>(
          loc, rewriter.getI32VectorAttr(sBoxAndMixColumnMatrix));
      SBoxAndMixColumnMatrix = rewriter.create<vector::BitCastOp>(
          loc, VectorType::get({272}, rewriter.getI8Type()),
          SBoxAndMixColumnMatrix);
      Value memSize = rewriter.create<arith::ConstantIndexOp>(loc, 272);
      Value sBoxBase =
          rewriter.create<memref::AllocOp>(loc, dynI8MemrefType, memSize);
      rewriter.create<vector::StoreOp>(loc, SBoxAndMixColumnMatrix, sBoxBase,
                                       c0);



      Value selCPU = rewriter.create<arith::ConstantIndexOp>(loc, i);
      Value isCPU = rewriter.create<arith::CmpIOp>(
          loc, arith::CmpIPredicate::eq, selCPU, cpuId);
      rewriter.create<scf::IfOp>(
          loc, isCPU, [&](OpBuilder &builder, Location loc) {
            Value lowerRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());
            sum += perGroupSum;
            Value upperRatio = rewriter.create<arith::ConstantFloatOp>(
                loc, APFloat(sum), rewriter.getF32Type());

            Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
            Value lowerI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), lowerF);
            Value lowerBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), lowerI);

            Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
            Value upperI = rewriter.create<arith::FPToSIOp>(
                loc, rewriter.getI32Type(), upperF);
            Value upperBound = rewriter.create<arith::IndexCastOp>(
                loc, rewriter.getIndexType(), upperI);
            if (i == groupNum - 1) {
              upperBound = numIdx;
            }
            Value groupNumIdx = rewriter.create<arith::SubIOp>(
                loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);
            Value groupSizeIn = rewriter.create<arith::MulIOp>(
                loc, /*lhs=*/groupNumIdx, /*rhs=*/msgLenIdx);

            Value byte_shift =
                rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound,
                                               /*rhs=*/msgLenIdx);
            Value groupMemIn = rewriter.create<memref::ViewOp>(
                loc, /*resultType0=*/dynI8MemrefType, /*source=*/candidates,
                /*byte_shift=*/
                byte_shift,
                /*sizes=*/ValueRange{groupSizeIn});

            auto ENV_printschedule = std::getenv("PRINT_SCHEDULE");
            if (ENV_printschedule && std::string(ENV_printschedule) == "ON")
              mlirPrintScheduleInfo(rewriter, op, i, groupMemIn, msgLen, 
                                    static_cast<int>(OP_TYPE::INPUT_MULTI_AES));

            // 'hyper.for' to run multiple no-dependency AES tasks
            auto hyperForOp = rewriter.create<hyper::ForOp>(
                loc, /*lowerBound=*/c0, /*upperBound=*/groupNumIdx,
                /*memIns=*/ValueRange{groupMemIn, targetOutput},
                /*memOuts=*/ValueRange{initMemref},
                /*odsArg4=*/
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange miargs, ValueRange moargs) {
                // result of the following scf::for, default: zeros
                auto for_result = builder.create<memref::AllocaOp>(loc, dynI8MemrefType, ValueRange{msgLenIdx});
                // set default value for for_result
                builder.create<scf::ForOp>(loc, c0, msgLenIdx, c1, ValueRange{},
                [&](OpBuilder &builder, Location loc, Value iv,
                    ValueRange) {
                  builder.create<memref::StoreOp>(
                      loc, /*value=*/c0i8, /*memref=*/for_result,
                      /*indices=*/ValueRange{iv});
                  builder.create<scf::YieldOp>(loc);
                });
                // traverse all candidate keys
                builder.create<scf::ForOp>(loc, c0, candidate_keys_count, c1, ValueRange{}, [&](OpBuilder& builder, Location loc, Value key_index, ValueRange args) {
                  // slice key from candidate keys
                  auto key_offset = builder.create<index::MulOp>(loc, key_index, candidates_key_len);
                  auto candidate_key = builder.create<memref::ViewOp>(loc, I8x16MemrefType, candidatesKeys, key_offset, ValueRange{});
                  Value expanded_key_size = builder.create<arith::ConstantIndexOp>(loc, 176);
                  Value keyBase = builder.create<memref::AllocOp>(loc, dynI8MemrefType, expanded_key_size);
                  VectorType vectorI8x16 = VectorType::get({16}, builder.getI8Type());
                  // for vectorize store
                  Value candidate_key_vec =
                      builder.create<vector::LoadOp>(loc, vectorI8x16, candidate_key, c0);
                  builder.create<vector::StoreOp>(loc, candidate_key_vec, keyBase, c0);
                  // do key expansion
                  keyExpansion(loc, builder, keyBase, sBoxBase);

                  Value inShift = builder.create<arith::MulIOp>(
                      loc, /*lhs=*/iv, /*rhs=*/msgLenIdx);
                  auto temp_output = builder.create<memref::AllocaOp>(loc, dynI8MemrefType, ValueRange{cptLenIdx});
                  Value input = builder.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/miargs[0], /*byte_shift=*/inShift,
                      /*sizes=*/ValueRange{msgLenIdx});
                  Value oriMsg = builder.create<memref::ViewOp>(
                      loc, /*resultType0=*/dynI8MemrefType,
                      /*source=*/miargs[0], /*byte_shift=*/inShift,
                      /*sizes=*/ValueRange{msgLenIdx});
                  // meaningless code, just for sst can get keyBase and sBoxBase
                  // which is out of the "hyper.for"
                  Value meaningless =
                      builder.create<memref::RankOp>(loc, /*memref=*/keyBase);
                  meaningless =
                      builder.create<memref::RankOp>(loc, /*memref=*/sBoxBase);
                  // AES computation
                  builder.create<crypto::AESOP>(
                      loc, /*input=*/input, /*msgLen=*/msgLen,
                      /*output=*/temp_output, /*key=*/keyBase,
                      /*sBoxBase=*/sBoxBase);
                  // compare
                  auto cond = builder.create<crypto::CmpOp>(loc, builder.getI1Type(), temp_output, miargs[1]);
                  // if matches, copy input to for_result
                  builder.create<scf::IfOp>(loc, cond, [&](OpBuilder &builder, Location loc) {
                    builder.create<scf::ForOp>(loc, c0, msgLenIdx, c1, ValueRange{}, [&](OpBuilder& builder, Location loc, Value iv, ValueRange) {
                      auto input_iv = builder.create<memref::LoadOp>(loc, input, ValueRange{iv});
                      builder.create<memref::StoreOp>(loc, input_iv, for_result, ValueRange{iv});
                      builder.create<scf::YieldOp>(loc);
                    });
                    builder.create<scf::YieldOp>(loc);
                  }, // scf::if
                  [&](OpBuilder &builder, Location loc) {
                      // else, do nothing
                      builder.create<scf::YieldOp>(loc);
                  }); // scf::else
                  builder.create<scf::YieldOp>(loc);
                }); // scf::for
                // reduce outputs
                builder.create<hyper::ReduceOp>(loc, ValueRange{for_result}, ArrayRef<mlir::function_ref<void(OpBuilder&, Location, Value, Value)>> {
                  [&](OpBuilder& builder, Location loc, Value lhs, Value rhs) {
                    builder.create<scf::ForOp>(loc, c0, msgLenIdx, c1, ValueRange{}, [&](OpBuilder& builder, Location loc, Value iv, ValueRange args) {
                      auto val = builder.create<memref::LoadOp>(loc, rhs, ValueRange{iv});
                      builder.create<memref::AtomicRMWOp>(loc, arith::AtomicRMWKindAttr::get(op->getContext(), arith::AtomicRMWKind::ori), val, lhs, ValueRange{iv});
                      builder.create<scf::YieldOp>(loc);
                    });
                    builder.create<hyper::ReduceReturnOp>(loc, lhs);
                  }
                });
                builder.create<hyper::YieldOp>(loc);
              });


            if (ENV_printschedule && std::string(ENV_printschedule) == "ON") {
              mlirPrintScheduleInfo(rewriter, op, i, candidatesKeys, cptLen, 
                                    static_cast<int>(OP_TYPE::EXTRA_MULTI_CRACK_AES));
              // TODO
              // mlirPrintScheduleInfo(rewriter, op, i, groupMemOut, cptLen, 
              //                       static_cast<int>(OP_TYPE::OUTPUT_MULTI_AES));
            }


            for (wafer::Chip c :
                 HardwareInfo.getArchitecture().getChips().getChipArray()) {
              float tmpRatio;
              if ((c.getType() == "CPU" || c.getType() == "GPU" || c.getType() == "CGRA") &&
                  c.getGroup() == i) {
                assert(perGroupSum != 0.0 && "exist lazy group!");
                tmpRatio = c.getDutyRatio() / perGroupSum;
                if (c.getType() == "CPU") {
                  // cpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else if (c.getType() == "CGRA") {
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("cgra")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                } else {
                  // gpu device dictionary
                  deviceSplit.push_back(rewriter.getDictionaryAttr(
                      {rewriter.getNamedAttr("targetId",
                                             rewriter.getStringAttr("gpu")),
                       rewriter.getNamedAttr("targetConfig",
                                             rewriter.getStringAttr("")),
                       rewriter.getNamedAttr(
                           "dutyRatio", rewriter.getF32FloatAttr(tmpRatio))}));
                }
              }
            } // end of chip split for-loop

            // add additional 'hyper.devices' attribute for 'hyper.for' op
            auto deviceSplitAttr = rewriter.getNamedAttr(
                "devices", rewriter.getArrayAttr(deviceSplit));
            auto isSharedMemAttr = rewriter.getNamedAttr(
                "isSharedMem", rewriter.getBoolArrayAttr({false, true, true}));

            SmallVector<NamedAttribute> attrs;
            for (auto attr : hyperForOp->getAttrs()) {
              if (attr.getName() != "devices" && 
                  attr.getName() != "isSharedMem")
                attrs.push_back(attr);
            }
            attrs.push_back(deviceSplitAttr);
            attrs.push_back(isSharedMemAttr);
            hyperForOp->setAttrs(attrs);

            rewriter.create<scf::YieldOp>(loc);
          });
    }

    // Determine if a cpu is a master cpu(0)
    Value isMasterCpu = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/cpuId,
        /*rhs=*/c0);

    // Determine if crack result is null
    Value Result0 = rewriter.create<memref::LoadOp>(loc, /*memref=*/initMemref,
                                                    /*indices=*/ValueRange{c0});
    Value isNullResult = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::eq, /*lhs=*/Result0,
        /*rhs=*/c0i8);

    // when cpu is a master cpu and crack result is null, cpu get crack result
    Value isMasterAndNull =
        rewriter.create<arith::AndIOp>(loc, isMasterCpu, isNullResult);
    rewriter.create<scf::IfOp>(
        loc, /*cond=*/isMasterAndNull,
        /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
          for (size_t i = 1; i < groupNum; i++) {
            Value srcid = rewriter.create<arith::ConstantIndexOp>(loc, i);
            rewriter.create<sst::MsgRecvOp>(loc, /*src=*/srcid, /*dst=*/c0,
                                            /*buffer=*/initMemref,
                                            /*size=*/msgLenIdx);
          }
          rewriter.create<scf::YieldOp>(loc);
        });

    // when cpu is not a master cpu and crack result is not null, cpu send crack
    // result Negate isMasterCpu and isNullResult
    Value notMasterCpu = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::ne, /*lhs=*/cpuId,
        /*rhs=*/c0);
    Value notNullResult = rewriter.create<arith::CmpIOp>(
        loc, /*predicate=*/arith::CmpIPredicate::ne, /*lhs=*/Result0,
        /*rhs=*/c0i8);
    Value notMasterAndnotNull =
        rewriter.create<arith::AndIOp>(loc, notMasterCpu, notNullResult);
    rewriter.create<scf::IfOp>(
        loc, /*cond=*/notMasterAndnotNull,
        /*thenBuilder=*/[&](OpBuilder &builder, Location loc) {
          rewriter.create<sst::MsgSendOp>(loc, /*src=*/cpuId, /*dst=*/c0,
                                          /*buffer=*/initMemref,
                                          /*size=*/msgLenIdx);
          rewriter.create<scf::YieldOp>(loc);
        });

    // replace the original operation with the result
    rewriter.replaceOp(op, ValueRange{initMemref});
    return success();
  }
};

} // namespace

void populateConvertCryptoToHyperPatterns(RewritePatternSet &patterns) {
  patterns.add<CryptoSha1MultiToHyperPattern, Sha1CrackLoweringPattern,
               CryptoMd5ToHyperPattern, MD5CrackLoweringPattern,
               CryptoSm3ToHyperPattern, SM3CrackLoweringPattern,
               CryptoAESToHyperPattern, AESCrackLoweringPattern_
               >(
      patterns.getContext());
}

namespace {
class LowerCryptoToHyperPass
    : public PassWrapper<LowerCryptoToHyperPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerCryptoToHyperPass)
  LowerCryptoToHyperPass() = default;
  LowerCryptoToHyperPass(const LowerCryptoToHyperPass &) {}

  StringRef getArgument() const final { return "lower-crypto-to-hyper"; }
  StringRef getDescription() const final {
    return "Lower Crypto Dialect to Hyper Dialect.";
  }

  void runOnOperation() override;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<func::FuncDialect, mlir::LLVM::LLVMDialect,
                    memref::MemRefDialect, vector::VectorDialect,
                    arith::ArithDialect, affine::AffineDialect, scf::SCFDialect,
                    gpu::GPUDialect, index::IndexDialect, hyper::HyperDialect,
                    sst::SSTDialect>();
  }
};
} // end anonymous namespace.

void LowerCryptoToHyperPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp module = getOperation();

  ConversionTarget target(*context);
  target.addLegalDialect<func::FuncDialect, vector::VectorDialect,
                         mlir::LLVM::LLVMDialect, memref::MemRefDialect,
                         arith::ArithDialect, affine::AffineDialect,
                         scf::SCFDialect, gpu::GPUDialect, index::IndexDialect,
                         hyper::HyperDialect, sst::SSTDialect>();
  target.addLegalOp<ModuleOp, func::FuncOp, func::ReturnOp,
                    crypto::PadMessageGetPadmem_OP, crypto::PadMessagePadone_OP,
                    crypto::Sha1core_OP, crypto::Sha1_OP, crypto::CmpOp,
                    crypto::MD5OP, crypto::SM3OP, crypto::AESOP>();

  RewritePatternSet patterns(context);
  populateConvertCryptoToHyperPatterns(patterns);

  if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();
}

void registerLowerCryptoToHyperPass() {
  PassRegistration<LowerCryptoToHyperPass>();
}

std::unique_ptr<Pass> createLowerCryptoToHyperPass() {
  return std::make_unique<LowerCryptoToHyperPass>();
}
